Generating Behaviorally Diverse Policies with Latent Diffusion Models

Summary

The paper presents a diffusion-based approach to compressing a policy archive discovered by a Quality-Diversity RL algorithm. The diffusion model operates in the latent space of a VAE and achieves high levels of compression together with a reasonable level of reconstruction. The algorithm is evaluated on a collection of 4 Brax environments.

Strengths

- Clearly written paper and description of the approach - Promising results showing high levels of compression on some Brax environments, together with good levels of reconstruction and coverage - Strong visualizations of behavior during training and different synthesized behaviors

Weaknesses

- The approach studied in the paper is limited to compression of the original archive. Equally, the sequential behavior composition experiments only reproduce what was possible with the original archive. An interesting next step would be understanding if the diffusion model can generalize to novel measure vectors or language instructions. - The paper assumes very small 2-layer, 128-width MLPs trained by the QD algorithm, it is unclear if the algorithm can scale to larger and more representative networks - High loss in diversity particularly on the ant environment, in Table 1. - Line 8 in the Abstract, should clarify exactly what environments the authors see the compression ratio/coverage on Minor: - Line 42: typo in ‘uspample’ - Scale is hard to see in Figure 3, axes should also be described - The idea of compressing policies into a single diffusion model is related to [1] which compresses offline RL datasets into a single diffusion model, [1] also achieves around 13x compression. - It would be helpful to also indicate the level of compression in Table 1. [1] Synthetic Experience Replay. Cong Lu, Philip J. Ball, Yee Whye Teh, Jack Parker-Holder.

Questions

- The ablations on network capacity in Table 2 are only run on humanoid which works well, it would be interesting to understand if the large loss in diversity on the ant environment can be alleviated by higher network capacity or better representation learned by the VAE. - Can the diffusion model generalize to novel conditions and improve on the original archive? - Section 3: The convolutional layers used to encode the weights and biases of the network could be less appropriate than equivariant layers such as those used in [1]. [1] Permutation Equivariant Neural Functionals. Allan Zhou, Kaien Yang, Kaylee Burns, Yiding Jiang, Samuel Sokota, J. Zico Kolter, Chelsea Finn.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

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

4 excellent

Contribution

3 good

Limitations

The limitations of the method are discussed well in the paper.

Summary

Quality-Diversity Reinforcement Learning generates a set of policies (here called the archive) that are learned to produce varied behaviors in the environment. These archives can be large, and this paper aims to compress a previously learned archive into a single model by leveraging a conditional diffusion model. Each of the policies in the archive is first represented in a latent space by using a variational auto-encoder that reconstructs the weights and biases of the policy's neural network. Once the encoder and decoder of the policies have been trained, a diffusion model is then fit to the encoded latents, and conditioning information is used to help guide sampling. The paper explores conditioning based on the measure of a policy (where the measure is a set of functions used to split the policies into different regions of behavior space), or text descriptions of the policies. The paper shows that the learned diffusion model can generate policies that return similar rewards as the original archive in aggregate, and also have high overlap with the conditioning measure when the sampled policy is executed again in the environment.

Strengths

This paper proposes an interesting application of powerful generative models to fit a Quality-Diversity archive of policies. The ability to reconstruct the full archive from a single generative model increases the practicality of QD approaches to skill discovery in reinforcement learning. The paper does a good job highlighting this contribution, and it is indeed an intriguing direction.

Weaknesses

- Evaluation of the model is thorough, in that ablations and several domains are explored, however it is difficult to assess the quality of the approach given that no alternative approaches are attempted. In line 210, the paper argues that other approaches to archive distillation are not comparable because the underlying archive is different. I disagree: since the main contribution of this paper is a distillation method, it should be able to compared to other distillation methods when the archive is held fixed. - I find the metrics used to evaluate the model difficult to interpret, perhaps related to the lack of baseline. The Mean-Error in Measure (MEM) metric described by the paper is a reasonable one, but the paper does not describe what measures are used in the various tasks, so interpreting the scale of MEM is difficult. QD score is similarly difficult to interpret. Is a decrease of $0.6 \times 10^7$ QD score a meaningful one? - The experiment on sequential behavior composition is similarly difficult to interpret. Is 80% success of 4 consecutive behaviors good? Perhaps including the success rate of the original archive would be a good start. - Graphs in Figure 3 are not immediately interpretable without prior exposure to that form of QD visualization.

Questions

- I find the design of the network encoder to be strange. Why are convolutional filters applied to the weight matrix of an MLP when encoding? Is there some spatial structure there to be captured? - What dependence does the diffusion model have on the underlying archive to be compressed? - Where do text descriptions of behaviors come from? Are policies labeled after learning, or are the labels pre-specified with the desired measures? Is it easy to create this conditioning data?

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed 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

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The authors have addressed the method's limitations.

Summary

This paper tried to solve the high space complexity in Quality Diversity and proposes a method that uses diffusion model to distill the archive into a single generative model based on policy parameters.

Strengths

* This paper leverages the generation power of diffusion model and condenses one model instead of thousands of policies. * This paper is well structured.

Weaknesses

* Quality Diversity is not well-known. It's better to include a background section instead of including basic knowledge in the related work. * It's confusing no baselines can be directly compared. The metrics can not only be rewards but also space efficiency.

Questions

* Need more explanation and comparison in experiment * There are some environments for quality diversity like QD-GYM. What's the performance of the method in QD-GYM?

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

2 fair

Presentation

3 good

Contribution

3 good

Limitations

N/A

Summary

This work presents a novel framework using latent diffusion models to distill the archive of polices into a single generative model over policy parameters. The latent diffusion model with VAE backbone compresses the high dimensional neural network (NN) parameters into a compact, making it possible to reconstruct policies parameterized by NN. Further, the conditioning mechanism of diffusion models is used to flexibly generate policies with different behaviors.

Strengths

1. The author(s) use latent diffusion model to compress the high dimensional neural network (NN) parameters into a compact, making it possible to directly generate policies parameterized by NN. 2. The proposed framework recovers 98% of the original rewards and 89% of the original coverage while achieving a compression ratio of 13x.

Weaknesses

1. The performance and accuracy of the proposed method show significant discrepancies when applied under text conditions. Experiments show that the success rate of the method is influenced by the selection of text labels. This suggests that the model possesses limited understanding and generation capabilities in terms of language descriptions. 2. The proposed method shows poor performance on tasks with high-dimensional measure vector, such as Ant, suggesting its limited modeling capability in high-dimensional measure spaces.

Questions

1. How many policies are used as training dataset when conducting archive reconstruction experiments? 2. Do the dimensions of the measures used in each environment have explicit semantic information? If so, can the author(s) further demonstrate the model's generalization ability on it? (For instance, assuming that a particular dimension of the metrics represents the robot's movement speed, if the value of that dimension is set to 1.2 in the condition of the diffusion model, will the generated strategies be able to achieve speeds beyond those observed in the training dataset?) If not, how could the users employ this model to generate policies of their desired behaviors?

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

1. The author(s) only demonstrated the model's ability to reconstruct the original dataset and did not conduct further experimental demonstrations regarding its generalization capability. However, the generalization ability should be an important consideration when evaluating generative models. 2. At present, it appears that text conditioning does not show a sufficiently favorable influence within the model.

Thank you

Thank you for the responses to the review. I will maintain my current score.

Thanks detailed response

I truly appreciate the author's detailed response. Thanks to that, all of my questions and concerns have been thoroughly clarified.