Paper: Policy Disentangled Variational Autoencoder (ICLR-2024)

Reviewer comment (Novelty):
The idea of introducing discrete modes in VAE for better future modeling is not new and has been explored in other fields at least 3 years ago. (e.g., Trajectron++ [3] (ECCV 2020) in motion forecasting)

Author response:
Thank you for taking the time to read and organize the review of our works. We appreciate your concerns about our proposed method and experiments. Following are the statements on the weakness of our paper. Please refer to the response to all reviewers for our response on contribution, baseline, and evaluation tasks.

The objective of our paper is not to generate high-quality video. Our argument and contribution are not to generate high-quality video while controlling the policy of an agent in the video. The contribution of the paper is that PDVAE categorizes the unlabeled video by the policy and generates video according to the policy. The experimental results with three datasets demonstrate the capability of PDVAE to learn the policy representation and controllability.

$\textbf{World model}$ utilizes the VAE architecture to generate images and video. However, the model does not have model architecture nor loss to learn and generate policy. The model learns the impact of action on the dynamics of the environment, overlooking the intention behind the action. We have raised the problem of overlooking in the introduction. $\textbf{Dreamer’s}$ objective is to learn the agent’s model dynamics. The model intends to explicitly learn the model dynamics. According to the objective of PDVAE, the model does not need to learn the dynamics explicitly. PDVAE learns the dynamics implicitly with the generative model. $\textbf{Trajectron++}$ learns the discrete latent feature $z_{mode}$ for better future modeling where PDVAE learns the discrete latent feature $\pi$ from the video and uses it to simulate diverse scenarios conditioned to the $\pi$. Trajectron++ does not control the generation with the latent mode. Considering the objective and methods used, we believe each model and PDVAE in different research criteria.

Thank you.

Reviewer comment (Experiments):
Despite the sophisticated architecture, there is little ablation study analyzing the importance of its various components.

Author response:
The model architecture of PDVAE might seem sophisticated, but essentially the model is composed of the policy-conditioned generative model (TD-VAE with policy), the policy extraction module, and the Siamese architecture for the regularization. We have done the ablation study on the impact of the policy-conditioned generative model on the policy extraction module with Moving MNIST and posted the qualitative results on the paper (Figures 6, 11, and 12). As we have illustrated in the second paragraph of Section 3.3, the Siamese architecture is essential to regulate the state. Without the structure, the conditioned policy does not affect the generation as the state contains complete information on the policy and environment.

Reviewer comment (Writing):
The "policy" in this work is really a set of learnable embeddings disentangled from the latent space, while the authors' description tends to confuse it with the concept of policy in reinforcement learning, which is a function / distribution over actions. I understand that the learnable embeddings are intended to capture the "policy" which hypothetically controls the "style" or "mode" of the video generation, but it would help readers understand if the authors could further clarify this and differentiate the distinct concepts.

Author response:
PDVAE uses the latent variable $\pi$ to encapsulate information on the agent’s intention or behavioral goal. Given samples of videos where each video contains an agent with its own strategy or behavioral goal, the videos can be categorized into discrete numbers by the strategies. From the perspective of the video generative model, such latent embedding can be interpreted as the “style” or “mode”. PDVAE implicitly learns the dynamics of the environment with TD-VAE-based architecture, so interpreting the latent variable as the policy is more appropriate. In reinforcement learning, the policy is a function to represent the agent’s behavioral intention. In PDVAE, we use the coded representation for the policy.

Reviewer comment (Experiments):
The paper builds on top of TDVAE. Why not compare with TDVAE in terms of video quality regardless that TDVAE cannot generate policy-conditioned rollouts?

Author response:
1. How does it compare with TDVAE, which the method is built on?

We have included the additional quantitative evaluations of TD-VAE with the same hyperparameters as PDVAE in Appendix F and below is the table with PDVAE. You can easily think of PDVAE without the policy module and Siam architecture in Figure 3 as the TD-VAE for the comparison.

As you have pointed out, we have thought that TD-VAE is not a valid baseline for the comparison because the contribution of PDVAE is to learn the representation of policy without supervision and the policy-conditioned generative model to control the generation.

Reviewer comment (Experiments):
The environments and tasks used in the evaluation are very toy cases where the action is always discrete. I wonder if the policy extraction module and the policy-conditioned dynamics model can deal with continuous actions.

Author response:
2. How does PDVAE work on real videos and continuous action space?

$\textbf{PDVAE does not explicitly learn the action space as the ELBO in Equation (6) indicates.}$ Our model is agnostic to the action space as it implicitly learns the action space. The common model architecture is used for all three datasets where Moving MNIST and VizDoom have discrete action space, and KTH has the continuous action space. In each transition, the agent moves a certain degree of arms and legs, not a discrete amount. Moreover, KTH is the real videos with low dimensionality. We have elaborated why we have chosen three dataset on the section 3 of common response.

Reviewer comment (Experiments):
How do you compare with other recent text2video approaches, e.g. Control-A-Video?

Author response:
3. How do you compare with modern text2video diffusion models? Could they solve the task in zero-shot?

We believe text2video diffusion model is not an appropriate comparison to our model. PDVAE and text2video diffusion models can serve similar functionality, but the size of the models and the objective of the model are different. We believe these are two different criteria. It would be great if you provide some more detailed comparisons between models.

Reviewer comment (Experiments):
How does it compare with TDVAE, which the method is built on?

Author response:
1. How does it compare with TDVAE, which the method is built on?

We have included the additional quantitative evaluations of TD-VAE with the same hyperparameters as PDVAE in Appendix F and below is the table with PDVAE. You can easily think of PDVAE without the policy module and Siam architecture in Figure 3 as the TD-VAE for the comparison.

As you have pointed out, we have thought that TD-VAE is not a valid baseline for the comparison because the contribution of PDVAE is to learn the representation of policy without supervision and the policy-conditioned generative model to control the generation.

Reviewer comment (Experiments):
How does PDVAE work on real videos and continuous action space?

Author response:
2. How does PDVAE work on real videos and continuous action space?

$\textbf{PDVAE does not explicitly learn the action space as the ELBO in Equation (6) indicates.}$ Our model is agnostic to the action space as it implicitly learns the action space. The common model architecture is used for all three datasets where Moving MNIST and VizDoom have discrete action space, and KTH has the continuous action space. In each transition, the agent moves a certain degree of arms and legs, not a discrete amount. Moreover, KTH is the real videos with low dimensionality. We have elaborated why we have chosen three dataset on the section 3 of common response.

Reviewer comment (Experiments):
How do you compare with modern text2video diffusion models?

Author response:
3. How do you compare with modern text2video diffusion models? Could they solve the task in zero-shot?

We believe text2video diffusion model is not an appropriate comparison to our model. PDVAE and text2video diffusion models can serve similar functionality, but the size of the models and the objective of the model are different. We believe these are two different criteria. It would be great if you provide some more detailed comparisons between models.

Reviewer comment (Experiments):
Could they solve the task in zero-shot?

Author response:
3. How do you compare with modern text2video diffusion models? Could they solve the task in zero-shot?

We believe text2video diffusion model is not an appropriate comparison to our model. PDVAE and text2video diffusion models can serve similar functionality, but the size of the models and the objective of the model are different. We believe these are two different criteria. It would be great if you provide some more detailed comparisons between models.