Paper: Learning Interactive Real-World Simulators (ICLR-2024)

Reviewer comment (Evaluation):
It would be nice if the paper delved more into the limitations of the models. The paper has shown that exciting results can be obtained, but it's useful for the community to know the limits of the generalization capabilities, especially if people want to use this in the future for various applications.

Author response:
Thank you for the positive feedback! We have expanded the discussion on the limitations of this work and included the following in the paper (conclusion section):

- Hallucination. When an action is unrealistic given the scene (e.g., “wash hands” is given to a tabletop robot), we observe hallucinations (e.g., the table turns into a sink or the view turning away from the tabletop robot and a sink shows up). Ideally, we want UniSim to detect actions that are not possible to simulate as opposed to Hallucinate unrealistic outcomes.
- Limited memory. UniSim conditioned on a few frames of the recent history cannot capture long-term memory (e.g., an apple in a drawer could disappear when the drawer is opened if putting the apple in the drawer is not a part of the history for conditioning). How much history to condition on depends on the application of UniSim (e.g., whether UniSim will be used for policy learning in a near-Markov setting or question answering that requires long-term memory).
- Limited out-of-domain generalization. This is especially true for domains that are not represented in the training data. For instance, UniSim is mostly trained on 4 robot morphologies, and its ability to generalize to an unseen robot is limited. Further scaling up training data could help, as UniSim’s training data is nowhere near all the video data available on the internet.
- Visual simulation only. UniSim is not suitable for environments where actions do not cause visual observation change (e.g., different forces in grasping a static cup). A true universal simulator should capture all aspects of the world beyond visual experience (e.g., sound, sensory, etc)

Reviewer comment (Reproducibility):
For reproducibility, it would be helpful if the authors could release the code and some example pre-trained checkpoints.

Author response:
We will work on releasing a code for doing inference with UniSim. We will try to open-source the UniSim model, but since the model is pretrained on private videos, we foresee challenges in this process.

Reviewer comment (Experiments):
One of the recurring claims throughout the paper is that the major novelty is UniSim's "orchestration of diverse datasets, each providing a different aspect of the overall experience" into a single model. Yet no hard evidence is given for this combination being important -- aside from two vague figures in Appendix E. At a minimum, it would be important to train a version of UniSim on, say, datasets from the Language Table environment _only_, and report numbers for when synthetic data was generated from this, in Table 2 and 3. This would help support the claim that dataset diversity is valuable.

Author response:
> 2. Effect of diverse datasets.

Thank you for pointing out the lack of dataset analysis in the current version of the paper. We conducted ablation studies on the effect of diverse datasets, and included the following table in Appendix E.1:
| Dataset | FVD | CLIP |
| ----------- | ----------- | ----------- |
| Internet only | 219.62 | 22.27 |
| UniSim without internet | 307.80 | 21.99 |
| UniSim   |   **211.30**     | **22.63** |

We found combining internet data with various human activity and robot data results in the best FVD and CLIP scores on a held-out test split, whereas only using internet data (last row in Table 5) or not including internet data (all except for last row in Table 5) leads to worse performance. We further provide more qualitative examples (https://drive.google.com/file/d/174qUvW8fUxvKBVv3UV4m2yZ8y3v6gMW6/view?usp=sharing) where training on entire UniSim data is better than using internet-only data or without internet data. In terms of using downstream RL and planning to evaluate the effect of diverse data, we trained a video model only on Language Table as suggested by the reviewer, but we did not observe significant performance differences between UniSim and this task-specific video model in Table 2 and Table 3. Our hypothesis is that Language Table is a simple domain and already has a lot of data, so training a domain-specific video model is sufficient, but this is not true in general for harder environments.

Reviewer comment (Experiments):
It would be useful to investigate how entwined the effect of actions is with the dataset distribution. For example, could camera commands (zoom in etc) successfully be applied to kitchen scenes as in Figure 3? The fact that the name of the dataset had to be included as part of the action during training, makes me suspect actions may not be able to generalise well to new kinds of video. This would not be a dealbreaker for the paper's acceptance, but is important to show readers how general this data mixing is.

Author response:
> 3.1 Entwined actions and videos.

We found that actions are generally not entwined with videos in the training distribution; we can apply diverse actions to scenes as long as such actions are reasonable. As the reviewer requested, we generate various egocentric movements in the kitchen scene (https://drive.google.com/file/d/1HhvJIOlA2TAotZiRrt0CHVkgjc0kEQH2/view?usp=sharing) from Figure 3, including look up, look down, zoom in, zoom out, etc. We see that the generated videos generally follow the actions, but the videos might look less realistic when the actions are less reasonable (e.g., zooming out from a person’s view looking down at the kitchen table). We also generated some diverse interactions with objects (https://drive.google.com/file/d/1MWzWOHpeC9etJsW6gZXHm2q2E8DKyR24/view?usp=sharing).

> 3.2 Dataset name as a part of action label.

Note that we only used dataset name for dataset with very few videos (habitat), as it helped the model identify the domain better. We think this may not be needed if we tune the data mixture ratio during pretraining. For all the other datasets, we don’t use dataset name, so that language actions from different datasets can generalize across domains.

Reviewer comment (Experiments):
A lack of strong baselines might be expected for this kind of scale-up work. But in their absence, ablations become more important, to verify that the various components of the model were all necessary. The paper only presents a brief study of which frames to condition on.

Author response:
> 4. Additional ablations.

Thank you for pointing out this weakness of the paper. In addition to adding the ablation of dataset diversity in Appendix E.1 and additional qualitative examples, we also conducted an ablation on the model size in Appendix E.2: 

| Model size | FVD | CLIP |
| ----------- | ----------- | ----------- |
| 500M | 277.85 | 22.08 |
| 1.6B | 224.61 | 22.27 |
| 5.6B   |   **211.30**     | **22.63** |

We found that while increasing the model size improves the video modeling performance, the amount of improvement measured by Fréchet Video Distance (FVD) plateaus as the model gets bigger, which is slightly disappointing from a scaling point of view.

Reviewer comment (Writing):
The model section is poorly written. The use of $\mathcal{T}$ is (I think) slightly misleading -- usually the transition fn of a POMDP is defined as operating on the states, $\mathcal{T}(s_t,a_t) \to s_{t+1}$, and there is a separate emission function producing the observations, $p(o_t|s_t)$. Eq. 1 implicitly combines these -- I might recommend renaming it $f$ or $g$ or whatever. I didn't follow why $o_l$ notation needed to be introduced, since it's immediately unrolled into $[o_t, o_{t+1}]$ and never referred to again. I also didn't understand why the model conditions on the noised, rather than clean, previous observations. It's said the last four frames are concatenated from $o_{t-1}$, which confused me -- does $o_{t-1}$ represent four frames, or should it read $o_{t-1:t-4}$ or similar?

Author response:
> 5.1 Confusion around the POMDP formulation.

Thank you for pointing out the confusion. We have reformulated UniSim as an observation prediction model $p(o_t | h_{t-1}, a_{t-1})$ in Section 2.2 (Page 4), where $h_{t-1}$ denoting the history up to time $t-1$ corresponds to the belief state in a POMDP. Under this formulation, sampling from $p(o_t | h_{t-1}, a_{t-1})$ corresponds to doing rollouts in the POMDP. We also provided motivation for this formulation (i.e., directly learning the POMDP requires learning distributions over $s_t$, which is difficult). We also clarified that in the POMDP of the real world, the frequency and duration of interactions vary (e.g., a human opens a door in 2 seconds, a motor executing a robot control in 0.2 second), so $o_t$ is a set of (variable length) frames and $a_t$ is a set of (variable length) control sequences or some high-level text actions. For data from most environments, we use the last four frames of $o_{t-1}$ for conditional generation of $o_t$. However, when $o_t$ only consists of a single frame, we condition on $o_{t-4},...,o_{t-1}$ to generate $o_t$, as both the last 4 frames of $o_{t-1}$ and $o_{t-4},...,o_{t-1}$ are specific instantiations of $h_{t-1}$.

> 5.2 Why the model conditions on the noised, rather than clean, previous observations.

The model is conditioned on the *clean* previous observations concatenated with the noise samples for the next set of observations. We made the wrong choice to call this concatenation the noisy previous observation, which has now been corrected.

Reviewer comment (Reproducibility):
It's a shame to give the model details only in the Appendix C, as I believe many readers would be interested in them. I hope some of these can be shifted to the main body, particularly key details around the diffusion architecture (such as the core and super-resolution modules) and the amount of compute required.

Author response:
> 6. Model details in main body as opposed to Appendix C.

We added a paragraph summarizing the architecture and training on Page 5 and point the rest of the architecture and training details to the appendix.

Reviewer comment (Experiments):
The two main experiments were conducted on environments that were within the training distribution of UniSim. It would have been more impressive to investigate the performance on new environments.

Author response:
> 8. Out-of-domain performance.

UniSim exhibits limited out-of-domain performance with both success and failure cases (https://drive.google.com/file/d/1tpInX4KUywXLt971Tr15_RsW8I_ytpvu/view?usp=sharing). We believe that training UniSim on more diverse robotic datasets (e.g., RT-X [1]) and scale up the data to more internet videos can potentially mitigate this limitation. We made note of this limitation in the conclusion section of the paper.

Reviewer comment (Presentation):
The wordy description of all datasets in 2.1, I felt was much better summarized by Table 5 in the Appendix (perhaps with the addition of a column explaining how an action space is defined and handled), and might be swapped. (Optional!)

Author response:
> 9. Wordy dataset description.

We made that section more concise highlighting different information covered in each dataset (the main point of that section). Thanks for the suggestion for including the dataset table in the main text, we will take this into account in final version.

Reviewer comment (Experiments):
Appendix says 1M steps on 512 TPUs with batchsize 256 -- this seemed a low ratio of training updates to available compute. Did performance saturate beyond this?

Author response:
> 10. 512 TPUs with batch size 256.

The model is large and we needed model parallelism, and hence the small batch size. The video modeling metrics converged with 1M steps.

Reviewer comment (Reproducibility):
What was the wall clock time of the model training?

Author response:
> Clock time and parameter.

UniSim has 5.6B parameters required 20 TPU days to train with 512 TPU-v3. We now report that in the paper (Page 5).

Reviewer comment (Reproducibility):
How many parameters were in the model?

Author response:
> Clock time and parameter.

UniSim has 5.6B parameters required 20 TPU days to train with 512 TPU-v3. We now report that in the paper (Page 5).

Reviewer comment (Reproducibility):
Will the model weights be open-sourced?

Author response:
> Open-source.

We will try to open-source, but since the model is pretrained on private videos, we foresee challenges in this process.