Solving Continual Offline RL through Selective Weights Activation on Aligned Spaces

We appreciate your valuable review and respond to your concerns below.

[I]. Explanation of weaknesses

[1/7] W1: It is really hard to understand the big picture of the article. Figure 1 is not really helping since it does not indicate data flow through the model. It is only by studying the algorithm in the appendix that I could understand what the various components actually do, in which order. That algorithm should be moved to the main text, or els Fig.1 should be improved and simplified.

We have revised Figure 1 to make the training and testing process of the model more specific and accurate. Regarding Figure 1, we should first focus on the QSA module in the bottom right corner. For the dataset of a specific task $1$, the QSA module is used to map the state and action spaces into a unified representation. Then, the top left section of Figure 1 illustrates the data flow during training and inference. Finally, the bottom left section of Figure 1 depicts the flow of data within the diffusion model.

[2/7] W2: This is not really continual learning, as you are processing all data in one go in Alg.1 during the QSA pretraining step. In a CL setting, you would have access to data from task i only when training on task i, but not before. So, from a CL perspective, adapting QSA in one go prior to real RL training is cheating because you are looking into the future. In the same vein: if you have access to all collected data, why use CL at all and not train everything at once?

Our method strictly follows the CL training paradigm. At any given training time, the network only has access to the dataset of a single task. Similarly, in the code, our method's training process also adheres to the CL training paradigm.

In the original pseudocode, we aimed to clearly express the hierarchical structure of the QSA module and the SWA module during the training of each task. Simultaneously, considering that the QSA module and the SWA module share the same "for each task "i" loop, the pseudocode was written in its original form. However, this has evidently caused some misunderstanding. We have revised the pseudocode to more accurately describe the entire training process. Below is the comparison of the original pseudocode and the new pseudocode.

The QSA module does not share parameters; all task-specific parameters, including the VQE, VQD, and task-related parts of the codebook, are independent. Therefore, during training, we can train the VQE, VQD, and codebook for each task sequentially. From the algorithm, we can also observe that the outermost loops for both the QSA module and the SWA module are "for each task" (lines 4 and lines 15). This makes it entirely feasible to train the QSA module sequentially for each task. To better illustrate how the QSA module aligns with the continual learning (CL) training setup, we have provided an additional Figure 15 to visualize the training process of the QSA module.

Original pseudocode: For each task i (lines 4):

• Train QSA module
• Save QSA module

For each task i (lines 15):

• Train SWA module
• Save SWA module

------------------------reformulate------------------------>

New pseudocode: For each task i:

• Train QSA module
• Save QSA module
• Train SWA module
• Save SWA module

[3/7] W3: It is not clear how this method could scale to more tasks when the mask is filled. It would be very helpful to comment on this, and why you do not simply train a separate DDPM for each task.

In continual learning, parameter isolation (i.e., structure-based) methods [1-4] we used are widely adopted to reduce catastrophic forgetting. Previous methods [5-8] aim to use a single model to gradually master solutions for multiple tasks following the CL training paradigm. In these approaches, the number of CL tasks is predefined, and they do not consider how to learn new tasks beyond the predefined CL task sequence. To address this challenge, we propose two solutions.

• Pruning Existing Masks that is suitable for a small number of new tasks (Refer to Figure 12 in our paper): Instead of increasing the number of masks, we prune the existing masks. Specifically, for each task, we identify less important weights, such as those with small absolute values, and release the corresponding masks for training in new tasks.
• Introducing New Masks that is suitable for large number of new tasks: In this approach, we keep the learned model parameters and masks unchanged and use additional parameters and masks to learn the new tasks.

I would like to express my sincere gratitude for taking the time to provide such detailed comments on our manuscript. Your thorough feedback is valuable and greatly appreciated. Accordingly, we included comprehensive theoretical analyses and experimental results to support our claims. These analyses demonstrate the robustness and validity of our proposed method. Furthermore, we kindly request that you reconsider the novelty and contributions of our work. Despite the existing research in this field, our approach introduces significant advancements and unique insights that address critical challenges and open new directions for future research. Thank you once again for your constructive feedback. We look forward to the opportunity to further discuss our work and address any remaining concerns you may have.

W7: again I disagree, the concept of offline reinforcement learning must be motivated since the essence of RL is in exploration. Some of the points you provide could make a case for offline RL, but the notion of offline continual RL is even more problematic: why do you need to learn tasks one after the other, if you could also learn them all at once in an offline fashion? As long as you cannot provide a convincing use case for this, I remain unconvinced, and I can hardly imagine what such a use case might be.

• Beyond "the essence of RL is in exploration" you claim, RL has so many other challenges, such as robustness, uncertainty, diversity, generalization, collaboration, long-term dependencies, sample efficiency, adversarial attack, credit assignment, ...
• No one can predict the kinds of tasks we may face in the future. So continual learning is proposed to continually master new skills to solve new tasks.
• In many scenarios, such as robotic control and autonomous driving, the commonly adopted approach is to train models on offline data. Additionally, with the rapid development of the generative model community, highly realistic synthetic data is becoming increasingly available. Training models on offline datasets will undoubtedly be a critical research direction.

Q1: Even if you argued that training can be done offline, at least testing must be done on-line. And again I find it quite hard to imagine a realistic scenario where someone would provide a task label at test time.

I will only say that your imagination is truly limited, to the point where I genuinely question whether you actually understand reinforcement learning. Take autonomous driving as an example: autonomous driving agents are trained using offline data and evaluated either in real-world environments or high-fidelity virtual environments. Additionally, robotics is another common example, where many studies are trained with offline data and then tested in online environments [6-8].

[1] A Definition of Continual Reinforcement Learning, Neurips

[2] Towards Continual Reinforcement Learning: A Review and Perspectives, JAIR

[3] Stable Continual Reinforcement Learning via Diffusion-based Trajectory Replay, ICLR workshop

[4] Diffusion-based Curriculum Reinforcement Learning, Neurips

[5] https://www.youtube.com/watch?v=NvfK1TkXmOQ

[6] D4RL: Datasets for Deep Data-driven Reinforcement Learning

[7] Continual World: A Robotic Benchmark For Continual Reinforcement Learning, Neurips

[8] Trafficgen: Learning to generate diverse and realistic traffic scenarios, TPAMI

W3: why not train a separate model for each task? It would be way easier and clearer, and probably more powerful, too.

Continual learning is proposed to address the potential and unpredictable demands of future tasks, making it highly relevant to real-world needs. Moreover, extensive research [1-5] has been devoted to tackling the challenge of catastrophic forgetting in continual learning.

• We must claim the advantages of training a single model for tasks with different state/action spaces: (1) It significantly expands the applicability of the model, eliminating the limitations imposed by task-specific state and action spaces. (2) It facilitates learning multiple tasks sequentially with a single model, removing the need to train a separate model for each task. (3) It can achieve higher computational efficiency and lower storage costs to master multiple tasks and make decisions with a single model.
• Let me give you an example: if you train one LLM as an agent and, after training, you want it to continuously acquire additional skills, according to your reasoning, we would need to retrain an entirely new LLM for every new skill. However, it is evident that this would result in enormous time and financial costs.

W5: just using official code is not enough. You need to document what parameter settings you used and where you adapted the code, if applicable. If you run the author's code with the authors' parameter settings, then results may be suboptimal because you did not adapt parameters to your problem. That is something that you need to do, too, otherwise the comparison is unfair.

Using official code and the default hyperparameter settings (precisely as reported in the original paper) is a consensus within the community; therefore, it is sufficient to demonstrate the effectiveness of our method.

• First, using the official code ensures the highest level of reliability in the results.
• Second, the default hyperparameter settings in the official code have already been meticulously tuned, making them the most representative of the baseline's performance.
• Third, reporting the default hyperparameter settings of baselines is unnecessary because we can already find the setting in their original papers.

W6: I strongly disagree, a single controller should be less efficient, whereas one controller per task would be the optimal case since this controller does not have to incorporate knowledge about other tasks. All in all, the necessity of having a single controller for problems that "live" in different state/action spaces seems very contrived to me, and the arguments you provide are not convincing to my mind.

Since you claim that "a single controller should be less efficient," why should we even pursue AGI in one agent? The reason we use a single model is to achieve higher computational efficiency and lower storage costs. Another example: current research on LLMs aims to use a single model to continuously learn and tackle most and even all tasks.

We are particularly encouraged that the reviewer finds our method novel and effective. We appreciate the valuable feedback and respond to your concerns below.

[I]. Explanation of weaknesses

[1/3] W1: Discussion of L2M and the experimental comparison of L2M and VQ-CD.

We discuss the differences between our method and the L2M method as follows:

• State Space Alignment: The L2M method aligns the state space using padding, where all environmental state vectors are padded to a length of 204 dimensions. In contrast, our method uses vector quantization to align the state space, resulting in an aligned dimension of 20, which is significantly smaller than the 204 dimensions used by L2M.
• Action Space Alignment: L2M treats each dimension of the action vector as a token and discretizes the action values with range [−1,1] into 64 bins. Using this approach, L2M can ignore the task-specific action dimensionality. For a task with an action dimension of m, L2M simply outputs m tokens using a for-loop to generate the m-dimensional action vector. In contrast, our method aligns the action space using vector quantization, where the action dimension is aligned to 10. We use an MLP network for this alignment, whereas L2M employs a transformer to output 64 dimensions.
• Comparison of Experimental Results: We conduct experiments on the D4RL dataset, where both the state and action spaces differ across tasks. We use the official implementation of L2M to conduct experiments [1]. The experimental results show that our method significantly outperforms L2M. We attribute this improvement to the following reasons: 1）Regarding L2M’s state padding approach, we have already included state padding baselines in our original submission (Figure 4), and the results demonstrate that our method outperforms the state padding baselines. 2）L2M’s action discretization approach divides the action range into bins, resulting in a minimum resolution of 2/64. In contrast, our method is not constrained by such limitations.

In addition to the comparisons above, we have provided further experimental comparisons in the table below. More experimental results can be found in the revised Figure 4 of the main body.

Table 1: The comparison of L2M and our method (VQ-CD).

[2/3] W2: Training needs to be conducted on each task separately, resulting in separate model checkpoints with different masked weights, this scales linearly with the number of tasks.

Although the Weights Assembling introduced in our paper is performed after the training is fully completed, it is evident that after finishing the training of each task $i$, we can obtain the corresponding mask $M_i$. As long as we have $M_i$, we can progressively save the network parameters without the need to store checkpoints from different time periods. The detailed process is as follows:

• 1. Construct an initialized model specifically for saving parameters, with the same parameter structure $W$ as the model being trained.
• 2. Train on the first task. Once training is completed, we obtain the checkpoint $W[i∗\omega]$ and the mask $M_1$.
• 3. Use $M_1$ to extract the relevant parameters for the first task from $W[i∗\omega]$ and save them to the corresponding positions in $W$. At this point, $W[i∗\omega]$ can be deleted, thus incurring no additional storage overhead.
• 4. Repeat steps 2 and 3 until training on all tasks is completed.

[3/3] W3: VQ-CD relies on task information, other methods [1] can be used in a task-agnostic manner.

Continual learning scenarios with explicit task boundaries and those without are two distinct continual learning settings, both of them are widely applicable in various domains [2-6]. Detecting task boundaries requires additional modules to compare task similarities. Admittedly, our method is not yet suitable for task-agnostic continual learning settings where task boundaries are unclear. However, we have successfully implemented continual learning across tasks with different state and action spaces in the explicit task boundary continual learning setting. In future research, we aim to explore the VQ-CD method further in task-agnostic continual learning settings.

Thank you for the additional changes!

The authors have addressed all my concerns and in my opinion substantially improved the paper. Therefore I am raising my score, as I believe this is good work and should be accepted.

[2/3] W2: This method is designed for task-aware settings with explicit task boundaries, which could reduce its applicability in task-free continual learning scenarios where such boundaries are not predefined.

Task-aware continual learning is quite prevalent in many applications, such as CW10 and CW20 continual learning settings [3], as explicit task boundaries can often be conveniently determined using various methods, such as leveraging large models or human feedback on existing datasets or constructing new datasets with explicit task boundaries based on classical RL algorithms [4].

Compared with task-aware continual learning, task boundary-agnostic continual learning requires an additional mechanism to detect whether a task change has occurred. Clearly, this demands extra task similarity measurement mechanisms for detection. Admittedly, this is currently a limitation of our approach. However, it is also one of the challenges we are actively addressing. We are confident that further progress will be reflected in our future research.

[3/3] W3: It seems that the propose method, through effective, is a combination of existing learning methods, e.g. VQ-VAE, Diffusion, tasks masks.

We must clarify that the main innovation of this paper lies in proposing a continual offline reinforcement learning method that can be applied to arbitrary state and action spaces by space alignment, which significantly broadens the scenarios where our method can be applied. Additionally, we systematically validate the effectiveness of our approach through extensive experiments.

Vector quantization is merely one way to achieve space alignment. Clearly, other methods, such as MLPs or padding, could also be used for this purpose. However, based on experimental comparisons, we found that vector quantization yielded the best results, which is why we adopted this approach for space alignment.

The application of diffusion models in reinforcement learning is already widespread, achieving highly competitive results on various datasets due to the models' strong expressiveness. This is the reason we chose diffusion models as the core structure for the continual learning component.

The masking technique is also one method used to realize parameter isolation in continual learning. We can also use a separate neural network for new tasks, such as LoRA. However, compared with LoRA, task-related weight masking is an effective method because we can hardly define the pre-trained model when applying LoRA.

It is important to emphasize that vector quantization and diffusion models with task-related masks are techniques to realize our idea of performing continual learning on space alignment rather than a simple combination of techniques.

[II]. Explanation of questions

[1/2] Q1: Can the learned masks be visualized? Visualizing the masks could reveal task similarities and therefore can provide important insight for this proposed method.

Yes, thank you for your suggestions. We select [Hopper-m,Walker2d-m,Halfcheetah-m] to visualize the weights mask in Figure 13, Appendix B.6 of the paper. To make it easy to show the mapping relation between masks and the weights, we draw the network structure and mask matrixes, where we only report the first 100 channels of the mask matrixes.

Merely relying on the mask matrix may not intuitively demonstrate the effectiveness of our method. Furthermore, we also visualize the aligned state features in Figure 14, Appendix B.7 of the paper. From the experimental results, we can see that the state features learned by the AE method are not well-mapped to separate regions but are instead mapped to multiple areas. In contrast, the features obtained by our method are better partitioned into individual regions, which is more conducive for the model to capture the data distribution.

[2/2] Q2: Can this method detect non-periodic task changes?

Detecting the task boundaries demands extra task similarity measurement mechanisms for detection. Admittedly, this is currently a limitation of our approach. It is one of the challenges we are actively addressing. We are confident that further progress will be reflected in our future research.

[1] Denoising diffusion implicit models

[2] Learning to Modulate pre-trained Models in RL

[3] Continual World: A Robotic Benchmark For Continual Reinforcement Learning

[4] Continual learning of large language models: A comprehensive survey

We are particularly encouraged that the reviewer finds our method novel and effective. We appreciate the valuable feedback and respond to the questions below.

[I]. Explanation of weaknesses

[1/3] W1: The proposed method is a little bit complicated and might require significant computational resources and engineering efforts.

In our paper, our method integrates the advantages of vector quantization with the strengths of diffusion models, enabling continual learning across tasks with different state and action spaces and achieving superior experimental performance. To ensure the sampling efficiency of the diffusion model, we employed diffusion model acceleration techniques, such as DDIM [1], reducing the generation steps from 200 to 10, resulting in a 19.76x speed improvement during experimental testing, as shown in Table 1. Furthermore, in Table 2, we provide a comparison between our method and the L2M method [2], including model parameters and GPU memory consumption. The experimental results also demonstrate that the computational overhead of our method is entirely acceptable.

Table 1: The comparison of generation speed with different generation steps under the CL setting of Ant-dir task-4-18-26-34-42-49. We conduct the experiment with NVIDIA GeForce RTX 3090 GPUs and Intel(R) Xeon(R) Gold 6230 CPU @ 2.10GHz. In the main body of our manuscript, we use the 10 diffusion steps setting for all experiments.

Table 2: The comparison of GPU memory consumption. We conduct the experiment with NVIDIA GeForce RTX 3090 GPUs and Intel(R) Xeon(R) Gold 6230 CPU @ 2.10GHz.

Regarding the engineering complexity, we made two simplifications to ensure the algorithm runs efficiently and to minimize the physical time required for training the model:

• We chose a dense optimizer instead of a sparse optimizer, primarily because dense optimizers are more efficient in updating parameters, requiring less time compared to sparse optimizers. Also, we conduct the physical time comparison of dense optimizer and sparse optimizer and report the resutls in Table 3, where the results illustrate the efficiency of training with dense optimizer and validate what we claim.
• In the paper, after the final training is completed, we apply Weights Assembling (introduced in Section 4.2) to extract the parameters corresponding to each task and combine them into a complete model. Rather than laboriously tagging which parameters need to be saved, this approach used in this paper (i.e., Weights Assembling) provides a more straightfoward and efficient solution.

Table 3: The comparison of time consumption per update between sparse and dense (normal) optimizers. We compare these two types of optimizers on the CL settings and find that when we first use the normal optimizer, such as Adam, to train the model and then use weights assembling to obtain the final model, the total physical time consumption is significantly smaller than sparse optimizer (e.g., sparse Adam).

We appreciate your valuable review and respond to your concerns below.

[I]. Explanation of weaknesses

[1/2] W1: The paper does not clarify how the 'return' is utilized. Based on Section 3.2, I guess that the algorithm employs the 'return' as a condition. So, is it the 'return' or the Q-function that is being used? Additionally, how is the target return or target function determined during the inference phase?

We will now elaborate on how the return is utilized during the training and inference phases, starting from the data processing stage. For aligned state and action features, we need to process the original state and action vectors with QSA module during training and inference. In the training phase:

• 1. For each trajectory in the dataset, assuming the trajectory length is $N$ , we calculate the cumulative discounted returns $\sum_{t'=t}^{N}\gamma^{t-t'} r_t$ from each timestep $t$ in the trajectory using a discount factor $\gamma$. Then we will obtain $\bar{R}_{1:N}$.
• 2. After the computation in step 1, we normalize the $N$ cumulative discounted returns to make the value range in [0, 1] and associate the normalized returns $R_{1:N}$ with the state at each timestep $\{s_t, R_t\}$.
• 3. Sample a batch from the dataset containing sequences of length $T_e$ consisting of states, actions, and their corresponding returns, i.e., $\{s_{t:t+T_e},a_{t:t+T_e},R_{t:t+T_e}\}$.
• 4. Use $\tau^0_s=s_{t:t+T_e}$ as the input to the diffusion model, and $R_{t:t+T_e}$ as the condition input for the diffusion model. Compute the diffusion model's loss and update the diffusion model parameters.
• 5. Use $s_{t:t+T_e-1}$ and $s_{t+1:t+T_e}$ as inputs to the inverse dynamics model, and use $a_{t:t+T_e-1}$ as the output targets to compute the inverse dynamics model's loss and update the model's parameters.

In the inference phase:

• 1. Receive the current state $s_t$ from the environment.
• 2. Sample a state sequence $\hat{s}_{t:t+T_e}$ with length T_e from the normal Gaussian distribution, replace the first state with $s_t$
• 3. Obtain the diffusion model's input $\tau^K_s=[s_t,\hat{s}_{t+1:t+T_e}]$.
• 4. Generally, following previous studies [1], we set the return condition $R=0.8$. Then, we can feed $\tau^K_s$ and $R$ into the diffusion model.
• 5. At each generation step we use diffusion model to calculate the noise $\epsilon_\theta(\tau^k_s, k, R)$.
• 6. Obtain $\tau^{k-1}_s$ by performing generation process with the diffusion model. After each generation step, we will replace the first state of generation output with $s_t$.
• 7. At the end of the generation process, we will obtain $\tau^0_s=[s_t,\bar{s}_{t+1:t+T_e}]$.
• 8. Finally, we use $s_t$ and $\bar{s}_{t+1}$ as inverse dynamic model's input to obtain action $\bar{a}_t$.
• 9. Interactive with the environment with $\bar{a}_t$.

[2/2] W2: Within the QSA module, which parameters are shared, and which need to be retrained for different tasks?

In the QSA module, there are no shared parameters. The primary purpose of the QSA module is to align the state and action spaces across different environments. Consequently, for different tasks, the internal components of the QSA module, vector quantized encoder (VQE), vector quantized decoder (VQD), and codebook are task-specific, and none of their parameters are shared. Thanks to the alignment provided by the QSA module, the inverse dynamics model in the SWA module can be shared. This is because the state and action spaces of different environments are mapped into an alignment space with the same value range.

[II]. Explanation of questions

[1/2] Q1: What is the "diffusion-based lifelong learning systems" in Abstract?

Diffusion-based lifelong learning systems refer to frameworks that utilize diffusion models for continual learning. Relevant research has been widely explored in the fields of computer vision (CV) and reinforcement learning (RL). Recent related algorithms include [2-6].

[2/2] Q2: Figure 1 appears somewhat cluttered and hard to understand. Labeling the sequence within the figure would enhance its clarity.

Thank you for your suggestions. We have labeled the state sequence in Figure 1 of the main body. Additionally, we have added the data flow for the loss computation of the diffusion model and the inverse dynamics model to better illustrate how the model is trained. If you have any further questions, we would be more than happy to address your concerns.

[1] Is Conditional Generative Modeling all you need for Decision-Making?

[2] Continual diffusion: Continual customization of text-to-image diffusion with c-lora

[3] Continual learning of diffusion models with generative distillation

[4] Continual Offline Reinforcement Learning via Diffusion-based Dual Generative Replay

[5] A Definition of Continual Reinforcement Learning

[6] Addressing Loss of Plasticity and Catastrophic Forgetting in Continual Learning