Inter-Environmental World Modeling for Continuous and Compositional Dynamics

Authors: Kohei Hayashi, Masanori Koyama, Julian Jorge Andrade Guerreiro
Primary area: learning on time series and dynamical systems
Keywords: world model, state space model, object centric learning, symmetry

Abstract:
Various world model frameworks are being developed today based on autoregressive frameworks that rely on discrete representations of actions and observations, and these frameworks are succeeding in constructing interactive generative models for the target environment of interest.  Meanwhile, humans demonstrate remarkable generalization abilities to combine experiences in multiple environments to mentally simulate and learn to control agents in diverse environments.  Inspired by this human capability, we introduce World modeling through Lie Action (WLA), an unsupervised framework that learns continuous latent action representations to simulate across environments. 
WLA learns a control interface with high controllability and predictive ability by simultaneously modeling the dynamics of multiple environments using Lie group theory and object-centric autoencoder. On benchmark synthetic and real-world datasets, we demonstrate that WLA can be trained using only video frames and, with minimal or no action labels, can quickly adapt to new environments with novel action sets.

Decision: reject
Reviewer scores: 5, 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1730058018138

thread_type:
official_review

<thread_object>
creator: reviewer_em9c

<content>
summary:
The paper introduces World Modeling through Lie Action (WLA), a generative framework that learns continuous, compositional actions for controlling agents across diverse environments using Lie group theory. Claimed contributions include:

Unified Simulator: A shared, environment-agnostic simulator leverages continuous actions, unlike models with discrete, environment-specific setups.
Controller Interface Problem (CIP): WLA solves CIP with adaptable controllers, even under minimal labeling.
Object-centric Modeling: Slot-based object-centric modeling supports interactions across multiple objects and actions.
Performance: WLA seems to outperform on benchmarks (Phyre and ProcGen), proving its generalizability and adaptability.
This approach enhances state-space modeling with flexible, multi-environment simulations

strengths:
The originality seems novel. The universal interface for control is an important problem to study. The emprical evidence seems to be strong, it can outperform the baseline significantly. I think there could be some contributions in the paper, but the presentation needs to be improved a lot.

weaknesses:
The paper’s ambitious framework for inter-environmental modeling has notable weaknesses, especially in mathematical clarity and formulation, which hinder its accessibility and rigor.
 Key Issues:
1. **Ambiguity in Transition Operator (\( g_{t, \delta} \))**: \( g_{t, \delta} \) maps specific observations \( x(t) \) to \( x(t+\delta) \), making it appear as a trivial one point to one point mapping. This definition does not capture the desired dynamic evolution. It does not make any sense in the form written in the paper. The paper acknowledges that \( g_{t, \delta} \) depends on individual trajectories but simplifies it as a generic operator, risking confusion. This omission detracts from the model’s mathematical precision, especially for multi-environment dynamics.

3. **Underdefined Action Space \( A \)**: The action space \( A \) and action \( a(t, \delta) \) lack clarity on whether actions are fixed or variable over intervals. This ambiguity in structure reduces the framework's comprehensibility in continuous control scenarios.

4. **Non-Compositional Transition Operators**: Due to the triviality of \( g_{t, \delta} \), composing transitions over time is not feasible, so why does it form a Lie group? The paper introduces Lie groups without a clear justification for their relevance to the specific environments. Additional reasoning would strengthen its argument for using Lie groups to model compositional and continuous dynamics.

To improve clarity, the paper would benefit from revisiting the definitions and dependencies in its transition operators, better defining the action space, and providing justification for its mathematical choices. More precise formulation and notation would significantly enhance its accessibility and application.

questions:
Revise the math definitions and formulations significantly and I will reconsider my score.

See weakness for all the math confusions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730547320637

thread_type:
official_review

<thread_object>
creator: reviewer_1ym7

<content>
summary:
In this paper, the authors present World modeling through Lie Action (WLA), a generative state-space-modeling framework that can be trained in an unsupervised fashion on interactive environments, and that allows controllable generation of future frames of a given environment. One main goal of this framework is learning continuous and compositional action representations, similar to what humans are able to do. It is a step towards an interactive world model that generalizes across environments, given these environments have a common basic rules of composition and continuity.
The paper presents a model built using this framework that is benchmarked on two datasets (ProcGen and Phyre). In general, the model can be trained using no or only few action labels, making it more versatile and faster to adapt to new settings with different action labels.

strengths:
1. The paper is clearly structured, and the individual components of the model with accompanying formulas, as well as the setting is well introduced. The authors carefully set the scope of the work and how it compares to similar approaches.
2. Benchmarks with different test settings (e.g. different FPS compared to train) shows the robustness of the model on this benchmark and the ability to infer continuous dynamics.
3. Given a main focus of the framework is the Lie group theory, I appreciate the ablation study that reveals the relevance of it (Rotation) to the predictive performance of the model

weaknesses:
My main concern is the limited amount of comparability to previous works and approaches. The model is only benchmarked on two datasets, and one is mainly used as a (successful) sanity check. The second dataset (with 16 quite diverse environments) is benchmarked against only one other model, though using several metrics. Even if it is a ‘first of its kind’ framework as stated by the authors, more benchmarks to confirm its proper function would increase trust in the method/framework. It can be acknowledged though that the number of available datasets/benchmarks for this specific scope is low, as also stated in a recent survey [1].
Other things that are not clear to me at the moment are in the question section.


[1] McCarthy, Robert, et al. "Towards Generalist Robot Learning from Internet Video: A Survey." arXiv preprint arXiv:2404.19664 (2024).

questions:
1. Was the encoder and decoder trained from scratch?
2. Did you do multiple runs with different random initializations, or are the results from just one experimental run?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730711788742

thread_type:
official_review

<thread_object>
creator: reviewer_vjwu

<content>
summary:
The paper proposes a method for enforcing a Lie group structure on the latent space of a multi-environment world model.  This structure imposes compositionality and continuity in the latent space, and allows the dynamics in the latent space to be linear.  The latent space is further structured by imposing object-specific dynamics (through the use of slot attention and assignment of "close" slots in nearby time steps).  This method allows to learn a common dynamic model from a sequence of observations sampled from a set of environments, and the paper shows that this model can generate observations that better reconstruct the ground truth observation compared to a strong baseline.

strengths:
- The paper presents an interesting idea, and a clever method for exploiting similarities between several environments at once.
- Learning reliable world models from observation trajectories without paired actions would be a dramatic step forward, and this paper presents a novel angle on that problem.
- Evaluation shows improvements over a strong baseline.

weaknesses:
- In places the presentation can be a bit unclear.  This is mostly in the description of the method.  I would have benefitted from a diagram of the steps required to train WLA and then use that to solve the CIP.  The current Figure 1 may not be necessary; the inter-environment aspect of WLA is probably the most straightforward step in the process.  
- Although not direct alternatives to the latent space presented here, there is some previous work in structured latent spaces that it may be interesting to compare against in the related work section.  eg.
    - Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images.  Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, Martin Riedmiller
    - Object Files and Schemata: Factorizing Declarative and Procedural Knowledge in Dynamical Systems.  Anirudh Goyal, Alex Lamb, Phanideep Gampa, Philippe Beaudoin, Sergey Levine, Charles Blundell, Yoshua Bengio, Michael Mozer
- Not a weakness, but Line 111 contains a reference to OCCAM, which I presume is a previous name for WLA?

questions:
- How are the parameters $N$ and $J$ chosen?  Do they need to scale with the complexity of the environment?  If so, how do we expect computation costs to scale? 
- Line 192 suggests an assumption about infinitesimal changes to observations between timesteps.  This assumption seems very strong, how would the method handle larger changes?
- How similar do the environments need to be for the inter-environment simulator to model them jointly?  For example, could you mix the Phyre and ProcGen environments?
- Does it matter how the dataset is sampled from each environment?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730715887775

thread_type:
official_review

<thread_object>
creator: reviewer_n9uv

<content>
summary:
The paper proposes a method akin to state-space models for simulating environments by applying Neural Fourier Transform. It uses Lie group theory to obtain a continuous representation of the environment dynamics.

strengths:
- The paper is clearly structured and well written
- The use of the Lie group theory and NFT to model the dynamics of interactive environments seems to be both novel and interesting
- The method is well motivated and theoretically grounded
- Good reproducibility (code provided)

weaknesses:
- No runtime analysis
- No discussion of hyperparameter selection
- No significance tests
- No error bars
- No scaling laws
- No quantitative results on Phyre

Unfortunately, the empirical evaluation does not meet scientific standards and, therefore, I do not deem the paper ready for publication yet. Once these problems are solved, I think this paper will constitute a nice contribution to the field.

questions:
- In Eq. (1), does the function Ctrl return an observation or a transition operator?
- In Eq. (3), shouldn't one take the limit of delta to 0 instead of t to infinity?
- Are there any limitations arising from the restriction of transitions to Lie groups?
- Since the latent space is divided into different slots for different objects and the slots forward dynamics are independent of each other, how can the model represent interactions between objects? 
- It seems that the size of the latent space is a crucial hyperparameter; how was it determined (both N and J)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732108615931

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response G1. Additional Experiment on the 1X World Model Dataset

comment:
To further validate our method, we compared WLA against Genie on the **1X World Model dataset**[^1]. This dataset comprises over 100 hours of videos capturing the actions of android-type robots in various environments and lighting conditions, including narrow hallways, spacious workbenches, and tabletops with multiple objects. We slightly adapted the architecture of our method to suit this setting, as detailed below, but otherwise followed the same experimental protocol as with the ProcGen dataset.

### **Results**

On the 1X World Model dataset, we observed results consistent with those from the ProcGen dataset. While Genie produces cleaner predictions for individual frames, it fails to generate video sequences that align with the provided action sequences. In contrast, our method successfully learns the correct responses, producing videos that more closely resemble the ground truth. 

These qualitative observations are supported by the quantitative results shown in the table below. We evaluated the models using three metrics:

1. **Peak Signal-to-Noise Ratio (PSNR)**
2. **Temporal Difference PSNR ($\Delta_t$PSNR)**
3. **Fréchet Video Distance (FVD)**[^2], utilizing the debiased version[^3]

Although our method slightly underperforms Genie in frame-wise evaluation (PSNR), it achieves better performance in temporally local evaluation ($\Delta_t$PSNR) and significantly outperforms Genie in FVD. This indicates that our generated videos are more natural and closer to the ground truth distribution.

We will include some of the predicted videos in the supplementary material.

#### **Quantitative Results on the 1X World Model Dataset**

|        | PSNR (↑)    | $\Delta_t$PSNR (↑) | FVD (↓)     |
|--------|-------------|---------------------|--------------|
| Genie  | **21.16**   | 0.78                | 393.85       |
| Ours   | 20.82       | **1.13**            | **131.02**   |

### **Conclusion**

These results demonstrate that WLA is capable of modeling real-world robot actions in 3D environments, in addition to 2D game environments. This enhancement addresses the concerns about limited experimental evaluation and increases trust in our method and framework.

### **Implementation Details**

In this experiment, each $256 \times 256$ RGB video frame was converted into an $8 \times 8$ grid of 18-bit binary tokens using MAGVIT2. Consequently, we trained our model on the time series of these tokens and replaced the reconstruction losses with logistic losses, as the tokens are binary. Additionally, since Slot Attention expects images as input, we employed a standard Vision Transformer (ViT) architecture to model the time evolution and alignment of the encoded tokens. Specifically, since there are $8 \times 8 = 64$ tokens in total, we treat them as if they are “slot” tokens. Also, we set the number of rotation angles to $J = 16$.

---

[^1]: Available at [https://huggingface.co/datasets/1x-technologies/worldmodel](https://huggingface.co/datasets/1x-technologies/worldmodel); we used version 1.1.

[^2]: Unterthiner et al., "Towards Accurate Generative Models of Video: A New Metric & Challenges," 2024.

[^3]: Ge et al., "Content-Debiased FVD for Evaluating Video Generation Models," 2024.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732108993164

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response G2. Hyperparameter Selection

comment:
Reviewers N9Uv and vJwu raised questions and concerns on the choice of hyperparameters, specifically the number of slots $N$ and the number of Lie group actions $J$. We agree that these are crucial and will include a detailed discussion in both the main text and the Appendix.

Generally, increasing $N$ and $J$ enhances the model capacity but also affects computational complexity, especially $N$, which directly impacts runtime. The same applies to $J$ since $D = 2J$.

We observed that the internal construction of $D$ significantly improves performance. Inspired by sinusoidal positional embeddings and G-NFT, we parameterized each rotation with a "frequency." We indexed $\theta$ with $i, j$ such that $\theta_{i, j}(t) = m_i \tilde{\theta}_j(t)$, where $\lbrace m_i \in \mathbb{R} \mid i = 1, \dots, F \rbrace$ are the rotation speeds and $j \in 1, \dots, \tilde{J}$ with $\tilde{J} = J / F$. We fixed $F$ with $D = 2J = \tilde{J} F$.

We conducted additional small-scale experiments on ProcGen to study the effect of varying $\tilde{J}$ and $N$. The results are summarized in the following table:

| $N$  | $\tilde{J}$ | Reconstruction Error (MSE) |
|------|-------------|----------------------------|
| 16   | 3           | 0.0583                     |
| 16   | 6           | 0.0587                     |
| 16   | 12          | 0.0439                     |
| 16   | 24          | 0.0378                     |
| 8    | 6           | 0.0641                     |
| 16   | 6           | 0.0587                     |
| 24   | 6           | 0.0397                     |

These results indicate that increasing $\tilde{J}$ with a smaller $F$ improves performance. Similarly, increasing $N$ reduces reconstruction error, highlighting the importance of these hyperparameters.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732619217260

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revised Paper Submitted

comment:
We have revised our paper based on your valuable comments. The main changes are as follows:

- **Added new experiments** (Section 6.3)
- **Improved figures** illustrating the model overview (Figures 1 and 2)
- **Enhanced mathematical rigor** (Sections 2, 3, 4, and Appendix A)
- **Included additional related work** (Section 5)
- **Added discussion on hyperparameters** (Section 4.4 and Appendix B4)
- **Corrected typos**

The major revisions are highlighted in blue in the manuscript.

We have also included videos of the new experiments in the supplementary materials. These videos clearly demonstrate the differences compared to the baseline methods, and we hope they provide additional support to the significance of our framework.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces World modeling through Lie Action (WLA), a method for learning action-conditioned video prediction models using an object-centric latent representation. It learns a continuous latent action space that generalises across environments.

Overall, the reviewers positively highlighted the clarity of the paper and the novelty of the theoretically-inspired framework.

There were some concerns about the validity of the theoretical assumptions made in the paper, but these were well addressed during the rebuttal. The main concerns raised by the reviewers, however, related to the quality and breadth of the experimental evaluation (simplistic tasks [PHYRE, ProcGen], single baseline).

The newly added robotics video experiments are a step in the right direction, but the evaluation setting considered is too short (in terms of predicted time horizon) to be very meaningful. For example, the video results in the appendix capture primarily static scenes with little to no motion over a very short time horizon that should be very easy to model. Given that these results were only added during the rebuttal period, it would be good to further refine these and get them properly reviewed at another conference.

To have a better chance of being considered at a top-tier conference like ICLR, the authors could, for example, also spend more effort on validating the method against an existing model on an existing (challenging) task to properly ground the performance of their method.