Diffusion-Based Offline RL for Improved Decision-Making in Augmented Single ARC Task

Weakness 1, 3 & Q3. It is true that relying on hard-coded rules to generate gold-standard trajectories makes augmentation more challenging for complex tasks. However, SOLAR includes augmented data for approximately 20 tasks, some of which are relatively complex. One of the primary goals of SOLAR is to ensure that the tasks learned by the model can be solved perfectly. Therefore, we expect that training with trajectories generated by SOLAR will demonstrate significant value in this regard. For simple tasks, we compared 100 test dataset problems with 500 training dataset problems and found no duplicates, which indicates that the model trained with SOLAR can solve new problems even within the same task.

While the augmentation process is not automated, solving ARC tasks inherently involves diverse and complex augmentation methods. ARC is a benchmark that inherently requires core knowledge priors and reasoning abilities to solve a diverse range of tasks. Solving ARC tasks in itself provides significant value, so we did not emphasize extending to environments beyond ARC. Rather than focusing on extending SOLAR itself to other benchmarks, we believe that advancing methodologies using SOLAR to solve ARC tasks could serve as a key approach for tackling other complex reasoning tasks in the future.

Weakness 2 & Q2. In this paper, LDCQ[1] was used as an example of an offline RL method to demonstrate that SOLAR can be used for learning through such offline RL approaches. Therefore, we did not conduct separate experiments on other baselines, such as non-diffusion policies. Instead, to demonstrate the effectiveness of the Q-function, we compared the sampling method (VAE prior, DDPM) using the VAE and Latent Diffusion Model used in LDCQ training with the RL method (LDCQ). As you mentioned, it would be worthwhile to compare with other methods. Since the main contribution is the introduction of SOLAR and demonstrating its applicability to offline RL, we have revised the main text to give greater emphasis to SOLAR's contribution. The revised sections are highlighted in blue, and we would appreciate it if you could review them.

Q1. We conducted additional experiments on two tasks beyond those used in the paper: 1) a task that performs different rules depending on the colors of two squares, and 2) a task where the edges of the grid are colored based on the given problem (Appendix C). The results showed that our method did not achieve high performance on either task. However, compared to the sampling method such as VAE, DDPM, using the Q-function as guidance led to better performance. Therefore, as mentioned in the Section 6 limitation, we believe that if the Q-function is improved, it could demonstrate reasoning abilities. Furthermore, solving ARC tasks in itself provides significant value, so we did not emphasize extending to environments beyond ARC.

Q4. Appendix A.3 provides information about the hardware used and the approximate time required for training. The amount of data required for training also varies significantly by task, and computational costs will differ accordingly. The variation in data requirements per task is a major challenge and may be a limitation of the current reinforcement learning methodologies. In future research, it may be useful to analyze the amount of data needed for each task, or consider a methodological shift to make the required data smaller and more consistent across tasks.

[1] Siddarth Venkatraman, Shivesh Khaitan, Ravi Tej Akella, John Dolan, Jeff Schneider, and Glen Berseth. Reasoning with Latent Diffusion in Offline Reinforcement Learning. In ICLR, 2024.

Weakness 1 & Q1. In this paper, LDCQ[1] was used as an example of an offline RL method to demonstrate that SOLAR can be used for learning through such offline RL approaches. Therefore, we did not conduct separate experiments on other baselines, such as non-diffusion policies. Instead, to demonstrate the effectiveness of the Q-function, we compared the sampling method (VAE prior, DDPM) using the VAE and Latent Diffusion Model used in LDCQ training with the RL method (LDCQ). As you mentioned, it would be worthwhile to compare with other methods. Since the main contribution is the introduction of SOLAR and demonstrating its applicability to offline RL, we have revised the main text to give greater emphasis to SOLAR's contribution. The revised sections are highlighted in blue, and we would appreciate it if you could review them.

Weakness 2. We conducted additional experiments on two tasks beyond those used in the paper: 1) a task that performs different rules depending on the colors of two squares, and 2) a task where the edges of the grid are colored based on the given problem (Appendix C). The results showed that our method did not achieve high performance on either task. However, compared to the sampling method such as VAE, DDPM, using the Q-function as guidance led to better performance. Therefore, as mentioned in the Section 6 limitation, we believe that if the Q-function is improved, it could demonstrate reasoning abilities.

Weakness 3. In this study, instead of comparing LDCQ with existing offline reinforcement learning methods, we adapted it to fit the unique characteristics of ARC tasks. The main focus of this research is to demonstrate that offline reinforcement learning can solve these tasks through a structured combination of actions that reflects the reasoning process required to solve ARC tasks. While it is also important to demonstrate that a representative offline RL method can learn effectively, the main contribution of this study lies in introducing SOLAR and attempting to solve ARC using diffusion-based RL. Additional experimental results are as mentioned in weakness 2.

[1] Siddarth Venkatraman, Shivesh Khaitan, Ravi Tej Akella, John Dolan, Jeff Schneider, and Glen Berseth. Reasoning with Latent Diffusion in Offline Reinforcement Learning. In ICLR, 2024.

Weakness 1, 4. In this paper, LDCQ [1] was used as an example of an offline RL method to demonstrate that SOLAR can be applied to offline RL. Instead of conducting experiments with other baselines, we adapted LDCQ to fit the unique characteristics of ARC tasks. To demonstrate the effectiveness of the Q-function, we compared sampling methods (VAE prior, DDPM) used in LDCQ training with the RL method (LDCQ).

The main contribution of this study lies in introducing SOLAR and demonstrating its applicability to offline RL. We have revised the main text to place greater emphasis on SOLAR's contribution, with the revised sections highlighted in blue for your review. The primary focus of this research was to show that ARC tasks can be solved using offline RL through action combinations, enabling reasoning and problem-solving.

Weakness 2 & Q2. To provide a more detailed explanation of the SOLAR generation method in Section 3.2, we have revised the sentences and added the algorithm for the SOLAR-Generator. The revised parts are highlighted in blue, and we kindly ask you to refer to Algorithm 1. More specific details and examples of the SOLAR generation process are extensively covered in Appendix B. Examples of the generated SOLAR can also be seen in Figure 4 of Section 4, which shows the SOLAR used in the experiments. And the detailed explanation of LDCQ that was disrupting the flow has been moved to Appendix A.

Weakness 3. We conducted additional experiments on two tasks beyond those used in the paper: 1) a task that performs different rules depending on the colors of two squares, and 2) a task where the edges of the grid are colored based on the given problem (Appendix C). The results showed that our method did not achieve high performance on either task. However, compared to the sampling method such as VAE, DDPM, using the Q-function as guidance led to better performance. Therefore, as mentioned in the Section 6 limitation, we believe that if the Q-function is improved, it could demonstrate reasoning abilities.

Q1. The advantage of LDCQ lies in its ability to compress multiple steps into a latent representation, predicting multiple steps at once. Based on ARCLE, most ARC tasks are solved within 5 to 10 steps, typically under 20, so we divided them into units of 5 steps—this is a tunable hyperparameter. In solving ARC, even for the same task, there are a wide variety of test examples, making it important to infer which stage of the problem-solving process the agent is in for unseen states during training. Therefore, latent representation can be beneficial. Additionally, while the ARCLE space may seem small due to its discrete nature, in practice, it is defined as a combination of operations and selections, making it more complex to predict than it initially appears.

Q3. Currently, the Q-function is computed in the form Q(s, a \mid \text{embedding}), where demonstration examples are embedded and incorporated into the calculation. $\boldsymbol{p}_{emb}$ denotes pair embedding. (Modify Appendix A Eq.3)

$$Q(s_t, z_{t}, p_{emb}) \leftarrow \left( r_{t:t+H} + \gamma^H Q\left(s_{t+H}, argmaxQ(s_{t+H}, z, p_{emb})\right), p_{emb} \right)$$

This serves two purposes. First, it uses the demonstration examples to classify tasks while learning the Q-function. Second, the actions may vary depending on the demonstration examples. As shown in Figure 12 (b), it is not possible to infer the action based solely on the state. Demonstration examples are required to accurately infer the correct color, and the demonstration pair must be considered to select the correct coloring action.

Q4. The explanation of LDCQ, as it is not a primary focus, has been moved to Appendix A. Figure 9 corresponds to the original Figure 3.

Although using TD learning is an important aspect of our approach, we did not include a separate illustration for it, as it only involves a minor modification of existing methodologies. The primary focus of this study is to demonstrate that the proposed reinforcement learning approach can effectively learn the desired behavior, which is why we developed SOLAR. Our emphasis lies in showcasing the viability of reinforcement learning for solving challenging reasoning tasks, supported by the structured data provided by SOLAR.

(b) illustrates the forward process of training the Diffusion Model. This follows the training process of the Diffusion Model used in the DDPM [2], where the model learns the forward process of adding noise to the original data. The denoising process, shown as (c) the backward process (DDPM sampling), is then used to generate data by reversing the added noise.

[1] Siddarth Venkatraman, Shivesh Khaitan, Ravi Tej Akella, John Dolan, Jeff Schneider, and Glen Berseth. Reasoning with Latent Diffusion in Offline Reinforcement Learning. In ICLR, 2024.

[2] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. In NeurIPS, 2020

Weakness 1. First, the augmentation methods used in [A1] and [A2] primarily involve simple transformations such as rotation, color changes, resizing, and cropping of the data. These methods are relatively straightforward compared to the approach employed in SOLAR. In the simpler tasks discussed in the paper, we applied similar methods—like adjusting size or color—which could also be considered augmentations. However, SOLAR goes a step further by augmenting not just a single task but various tasks, allowing for the creation of new states by modifying elements such as the number of objects or their arrangement within the entire state grid.

In other words, SOLAR is capable of not only making simple transformations to existing states but also modifying the rules that apply across all demonstration and test examples for a given task. It even creates new states that adhere to these modified rules. For instance, as shown in Figure 9 of Appendix B3, which is based on Task 2 from Figure 1 in the Introduction, we modified the rules associated with each color and generated new positions for the 2x2 squares. Users can even adjust the size of these squares if needed. Therefore, this process generates new problems that share the same underlying rules, allowing for diverse problem generation within a single task. This enables AI systems to learn the analogy required by the task and assess whether the learned analogy can be effectively applied.

In comparison, the RE-ARC dataset provides a method for augmenting input-output pairs. However, as highlighted in Appendix C, the key distinction of SOLAR lies in its ability to augment intermediate states during the solution process of ARC tasks. SOLAR incorporates transitions derived from reinforcement learning into the generated trajectories. For instance, as illustrated in Figure 9, the SOLAR-Generator can adjust object colors, rearrange them, and randomly redefine the meanings associated with each color.

Each ARC task has specific rules that must be followed to solve it. To claim that an AI system has fully understood the rules of a task, it should not only solve a single instance but also demonstrate strong performance across a variety of problems that share those rules. Therefore, we believe that it is essential to train trajectories across diverse input-output pairs to evaluate understanding effectively. Ultimately, the goal of SOLAR is to provide a dataset that allows ARC tasks to be tackled using offline reinforcement learning methods, while also generating diverse problems to evaluate whether the rules of the task have been fully understood.

Weakness 2. We understood 'overshoot' as a situation where the agent goes beyond the desired correct state.

In ARCLE, rewards are only provided when the agent executes the $`Submit`$ action at a correct state. This ensures that any episode with overshoot is inherently treated as non-optimal. Gold-standard trajectories reach the correct state in the most efficient manner, while non-optimal trajectories may or may not reach the correct state. When generating gold-standard trajectories, the SOLAR-Generator verifies that the final state of the trajectory is indeed correct. In the ARCLE-based environment, an agent only receives rewards if it executes the $`Submit`$ action at the correct state. Therefore, episodes that overshoot the correct state are automatically classified as non-optimal.

Even in situations where non-optimal episodes are mixed with gold-standard trajectories, our experiments demonstrate the agent’s ability to follow the gold-standard trajectory. As shown in the results in Figure 7(a), the metric for "reach answer" is higher than that for "submit answer." This indicates the presence of test episodes with overshoot, where the agent reaches the correct state but fails to execute the $`Submit`$ action at the appropriate step.

Weakness 3 & Q3. Not all trajectories are generated from the same input-output pair. The task used in the experiment is named 'Simple task,' and the precise augmentation process for generating trajectories is as follows:

1. Input-Output Pair Generation: Grid Maker creates new input-output pairs according to the conditions of the given task. In this experiment, 500 test example input-output pairs were generated.
2. Action Sequence Generation: Grid Maker generates the action sequence required to solve the test example. Each task's Grid Maker has a hard-coded algorithm that calculates which operations and selections are needed to produce the output from the generated input. The resulting action sequences are referred to as 'gold-standard' because they solve the problems without any unnecessary actions. From a certain point in the gold-standard, random actions are added to create 'non-optimal' trajectories. In the Simple task used in the main experiment, eight random actions are taken from the branch point, followed by Submit. Through this process, a total of nine non-optimal action sequences are generated from each gold-standard action sequence. Most non-optimal sequences do not reach the correct answer, but some reach it by chance, contributing to the diversity of the solution trajectories.
3. Dataset Composition: Based on the generated input-output pairs and action sequences, the dataset is composed by obtaining information such as intermediate states, rewards, and termination states through ARCLE.

Additionally, we checked for overlap between the evaluation set and training set test examples. In the Simple task, the evaluation set was composed entirely of new test examples, allowing us to verify that the model understood and applied the task rules even to previously unseen examples. This demonstrates the model's generalization capability. While the states in non-optimal sequences may not be necessary for the Simple task, they could be valuable for other tasks in a multi-task environment, potentially enhancing the model's generalization ability. We also modified the content of section 4 to better understand.

Weakness 4. In this paper, LDCQ[1] was used as an example of an offline RL method to demonstrate that SOLAR can be used for learning through such offline RL approaches. Therefore, we did not conduct separate experiments on other baselines, such as non-diffusion policies. Instead, to demonstrate the effectiveness of the Q-function, we compared the sampling method (VAE prior, DDPM) using the VAE and Latent Diffusion Model used in LDCQ training with the RL method (LDCQ). As you mentioned, it would be worthwhile to compare with other methods. Since the main contribution is the introduction of SOLAR and demonstrating its applicability to offline RL, we have revised the main text to give greater emphasis to SOLAR's contribution. The revised sections are highlighted in blue, and we would appreciate it if you could review them.

Q1. As mentioned in the response to Weakness 3, multiple non-optimal trajectories are generated based on a single gold-standard trajectory for each input-output pair. Each task's Grid Maker has a hard-coded algorithm that calculates the necessary operations and selections to produce the output for the given input. The trajectory generated through this algorithm is the gold-standard trajectory. Since gold-standard trajectories are verified to ensure they successfully reach the correct state, they always achieve the correct answer. In contrast, non-optimal trajectories are intentionally included to lower the overall quality of the dataset and increase the decision-making difficulty. We aimed to demonstrate the potential of reinforcement learning by finding actions that contribute to solving the problem within a dataset containing unnecessary actions from non-optimal trajectories.

Q2. The original ARC training data includes several demonstration examples but only one test example per task, which is insufficient for training deep learning models effectively. We believe that a model that can solve a diverse set of test examples, all considered part of the same task, is truly capable of understanding and solving that task. Our goal is to train the model by providing sufficient data for tasks that share the same rules, thereby enabling it to fully understand and reason through the given task. Therefore, rather than focusing on performance improvements using an augmented dataset compared to the original ARC data, this study aims to highlight the importance of enabling sequential decision-making and reasoning through offline reinforcement learning, which would be difficult to achieve with the original ARC data.

[1] Siddarth Venkatraman, Shivesh Khaitan, Ravi Tej Akella, John Dolan, Jeff Schneider, and Glen Berseth. Reasoning with Latent Diffusion in Offline Reinforcement Learning. In ICLR, 2024.