Learning Dynamic Attribute-factored World Models for Efficient Multi-object Reinforcement Learning

Summary

This paper proposes a new relational world model which is at an attribute level. Based on the state knowledge from the environment or unsupervised object-centric representation (OCR) learning [1,2], they collected attribute-level object-centric knowledge, thereby, based on them, they learned a relational world model. In the aspect that their model can infer the detailed relationship between action and the specific attributes of objects (e.g., only change on the object position through “push” action, not changing the object’s color or shape) or the objects, their world model is more compositional and can be well generalized. They evaluated it through diverse generalization experiments on diverse three benchmarks, some of which give the ground-truth state information and others only give visual inputs. [1] SM Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey E Hinton, et al. Attend, infer, repeat: Fast scene understanding with generative models. Advances in neural information processing systems, 29, 2016. [2] Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances in Neural Information Processing Systems, 33:11525–11538, 2020.

Strengths

- This paper proposes a new attribute-level factorized world model for multi-object RL tasks. It is an excellent approach to step deeper from an object-wise world model such as C-SWM [1]. Through this, they can extend the compositionality to the attribute level. - They well discussed which model they designed and how to learn the learnable parameters to fit. - They evaluated their modeling with diverse world model methods and prior works which studied policy learning with unsupervised object-centric representation learning. - Through their evaluation, they can find their model outperforms baselines, significantly, when generalization is more required. - They also studied ablations to analyze which components are more important for the performance. [1] Kipf, Thomas, Elise Van der Pol, and Max Welling. "Contrastive learning of structured world models." arXiv preprint arXiv:1911.12247 (2019).

Weaknesses

- Their model requires the class label to learn the class template graph through multiple objects in the class. It is reasonable to learn the template graph per class and it makes sense to reduce the complexity of the modeling and improve the performance by merging to learn a graph with multiple objects in the same class. However, it is an unrealistic assumption for real-world data. For example, if we applied this model to real-world data, then we should require additional labor to do labeling for each object. - Their model is assumed that the attribute can be accessible. Similar to the class label, it could be hard to be assumed now. For example, for AIR [1], it could be reasonable, because the AIR encoder returns position and size separately, but it is not for recent OCR models such as SA [2] or SLATE [3]. Recently, a few works have tried to represent the OCR through disentangled representation vectors [4], so we can expect that the attribute (maybe) can be given through the encoders. It is curious how you applied SA to your model. - Defining the reward function as an object level is too limited to cover general tasks. Sometimes the reward could be defined through the relationship between objects such as the object comparison task in the Image benchmark. It is curious how you made your model solve the task. In your modeling, reward estimation is learned through step 1 (Class Learning in single-object environments). - Action binding is unclearly discussed. In the section for DAFT-MDP, you assumed that an action has only an effect on one object at a time (line 74). However, action binding can affect multiple objects through soft attention (section 3.2.1), and I cannot find the details to regulate the action effect to only a single object from the detailed criteria (line 690 in Appendix). - Fine-tuning to a new multi-object domain is unclearly discussed. Which parameters are fine-tuned? Does Fine-tuning mean the parameters in the world model are frozen and only learning the parameters for the policy? - (minor) typos - In line 222, $i$ should be $m$. - In line 235, $o_i^t$ should be $o_1^t$. - In line 236, isn't the value $v^t = <f_v(o_1^t),...,f_v(o_m^t)>$? - In line 257, the input, isn't it $h^{t+1}_{(i,j)}$? - PPO is not MBRL, but it is mentioned as one of them in line 293. [1] SM Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey E Hinton, et al. Attend, infer, repeat: Fast scene understanding with generative models. Advances in neural information processing systems, 29, 2016. [2] Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances in Neural Information Processing Systems, 33:11525–11538, 2020. [3] Singh, Gautam, Fei Deng, and Sungjin Ahn. "Illiterate dall-e learns to compose." arXiv preprint arXiv:2110.11405 (2021). [4] Singh, Gautam, Yeongbin Kim, and Sungjin Ahn. "Neural systematic binder." The Eleventh International Conference on Learning Representations. 2022.

Questions

- In the ablation study (D.3.2 in Appendix), how did you represent the interaction except for the class template or dynamic interaction graphs? For example, when the action is to push the object, then how is the action represented? Can we see the ablation without the class template graph that is under the assumption that every object is in a single class? - What is TFD in line 852? - In [1], the Transformer encoder (similar to the self-attention that you evaluated, but using a learnable token for policy) shows good performances for the Image benchmark. Could you compare it with your model? - You utilized AIR [2] for the approaches relying on symbolic inputs. Could you test them with SA [3], not AIR? A few tasks are enough to evaluate due to the time limitation for the rebuttal. [1] Jaesik Yoon, Yi-Fu Wu, Heechul Bae, and Sungjin Ahn. An investigation into pre-training object-centric representations for reinforcement learning. International Conference on Machine Learning (ICML), 2023. [2] SM Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey E Hinton, et al. Attend, infer, repeat: Fast scene understanding with generative models. Advances in neural information processing systems, 29, 2016. [3] Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances in Neural Information Processing Systems, 33:11525–11538, 2020.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, 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

The authors didn't discuss their model's limitations. As I roughly discuss it, their modeling requires a class template that can restrict the application of their modeling, not just for collecting the class label, but also for their assumption for the class that in the class, the attributes and interaction patterns are shared. Another limitation is that it requires the encoder or environment must give the attribute-level representations. Because their model utilizes them as an input and the target of their objective, one-step prediction. Except for those limitations on their design, there is no ethical issue with their work I think, because it is fundamental research work rather than applications that can affect the real world. Additionally, they didn't discuss the computational overhead of their training or progressing their work.

Summary

This paper proposes a methodology to tackle compositional generalization, allowing the policy to be able to generalize to previous unseen combinations of objects or compose previously learned tasks. More specifically, the authors propose the Dynamic Attribute Factored RL (DAFT-RL) framework, which involves learning various graphs which model the objects and their interactions. Experimental results reveal superior performance relative to many baseline methods from the literature.

Strengths

- Well-written. The paper is well-written, and the figures aid in the understanding of the methodology. - Strong experimental results. The proposed method is compared to a number of methods from the literature and outperforms all of them. - Ablations for major components. The paper includes ablations for the major components of the design, showing that all of the major components help in performance.

Weaknesses

- Comparison to prior work. It is unclear whether the comparison to prior work is completely fair. For instance, an imagination component is added to the baselines -- could this hurt the performance of the baselines? Furthermore, Table A4 in the Appendix shows strong results for some of the baselines (SMORL, NCS) during training, outperforming the proposed methodology, potentially implying overfitting. Is the capacity of the baselines comparable to that of the proposed method? Finally, is it possible that by using SA / AIR or the ground truth latent vector z from the simulator for image-based methods, you are actually hurting performance of these methods which use a visual encoder other than SA / AIR? - Choice of tasks. The best performing methods apart from the proposed method (NCS and LRN) evaluate their approaches on a different set of tasks to the one used in this work. It is unclear why the authors chose to deviate from this choice of tasks. NCS [12] may be the SOTA and the most recently published, so it would be beneficial to compare to NCS under their settings. - Incomprehensive ablations. The proposed method seems to be rather complex, and it is not convincing that such a complex method is needed. While ablations have been provided for the utility of each component at a higher level of abstraction, more ablations on the design choices for each module (factored class template graph, dynamic interaction graph, factored interaction pattern) would be beneficial. For instance, how critical is the action binding in the dynamic interaction graph? How important are the design choices in the dynamic action binding network, e.g. the choice of using soft-attention networks?

Questions

- The lack of the dynamic interaction graph leads to the biggest drop in performance in the ablations -- If there are two blocks and two switches, and the number of interactions is limited, why does modeling the dynamic interactions in a graph help so much? - How would an end-to-end approach do?

Rating

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

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central 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

3 good

Contribution

3 good

Limitations

There is no discussion about the limitations of the work.

Summary

The authors propose DAFT-RL, which uses object-centric representations to improve generalization to object attributes, when training world-models for model-based RL. For each class of object they learn a class template graph that maps class attributes to dynamics and rewards. The authors claim that this representation outperforms state-of-the-art generalizing to unseen objects with new attributes on a simulated block stacking task. FYI: Supplementary Material PDF just links to the paper pdf for me. Maybe intentional?

Strengths

- The object-centric attribute factorization is a reasonable approach when the number of attributes that affect environment behavior (dynamics and rewards) is small, or when there are multiple objects that share attributes. - On the benchmarks shown, the model does outperform baselines (although more on this in weaknesses). - I found the explosion easy to follow and the model section is well written (there's a fair bit of notation but at least it's clear).

Weaknesses

- "e.g. an object’s position and reward are affected by its previous position, but not by its appearance or activation state" - nit: appearance and dynamics are not independent in the real world. A metal object appears metal. A rubber object appears rubber. One can infer a lot about dynamics from appearance. I get the authors point, but consider rephrasing for clarity. - "Often an object’s transition and reward functions" - Agreed but a reward for a given object is also task dependent. "Stack the red block" means the "blue block" has no reward (or even a penalty). Reword for clarity. - At a high level, the overall method is really just an extension to Object-Oriented (PO)MDPs to include object attributes. I'm not sure this crosses the bar for me for technical contribution. Adding attributes seems like an implementation detail that will help in certain applications, rather than a fundamental algorithmic contribution. - Experimental evidence is probably the weakest component. OpenAI Fetch, SpriteWorld and Block Stacking are all very simple simulated environments. I think the world-model RL folks have moved passed these and they stand more as unit tests than real RL problems. Real-world results would be nice, but not necessarily required. In particular I can see this method failing when the factored attribute graph is hard to infer or ambiguous (think - a "rubber duck" that is just painted metal to look rubbery). I don't see how these simulated environments is really stressing the generative model. - Likewise, the environments were largely chosen to result in the baselines not working, so I'm not entirely sure seeing margin there is all that compelling. i.e. choose a set of environments with varying attributes until baselines fail, then explicitly predict those attributes so that our model works, doesn't seem overly compelling as an experimental platform. If it was on a real-world experimental domain (with all the associated complexity of real data), that would be more compelling. - I also worry about how far this general approach scales. Assuming a known set of classes with known observable attributes isn't going to work for complex scenes. Likewise, how how can you generalize to a unknown attribute that is not within your set? The rigidity of the object classes and attributes feels perfectly tailored to simple simulated environments.

Questions

None really. I didn't find any parts of the paper confusing.

Rating

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

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

3 good

Contribution

1 poor

Limitations

Is there a limitations section? The "modelling assumptions" is perhaps closely related. There's also no societal impact statements, although I don't think it's necessary in this case given the scope of the experimental results (toy simulations).

Summary

The paper presents a new framework for learning world models of environments with multi-object interactions. The presented DAFT-RL framework identifies objects as instances of classes, for which class template graphs describe how the object's dynamics and reward depend on it's attributes, and interaction pattern graphs describe how instances of different classes influence each other in terms of their attributes. The manuscript also presents how to learn a policy from this graph-based world model and shows that this policy can be adapted without additional policy learning to novel configurations by inference of object interactions and latent parameters. The proposed method is evaluated on several datasets and is shown to yield strong performance.

Strengths

* To the best of my knowledge the proposed method is novel * Design decisions in DAFT-RL are well-motivated and the overall idea is well explained. * The performance improvements are impressive, especially in the environments with more objects.

Weaknesses

1. I think it would be a lot easier for readers to parse the examples in Figure 1 if they were annotated with example attributes and parameters, e.g. position, velocity, friction coefficient, etc. 1. I didn't see a description of how the hyper-parameters were selected.

Questions

3. Line 66: I think it should say "Objects of the same class will have the same **set of**" to be as explicit as possible. I had to reread the sentence a couple of times. 1. Line 121: "some of the attributes of each object after the other". I think it should be "affect" instead of "after" 1. Line 797-798 (Appendix D.2.2): Is that two layers before and two after, or two, one of which before, the other after? ## Acknowledgement of rebuttal I have read the rebuttal and other reviews. The rebuttal addressed my concerns by describing how hyperparameters were selected, providing a discussion of limitations, and some ideas on how to address these limitations. The rebuttal has convinced me of the appropriateness of my accept rating.

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

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

I did not see a discussion of limitations. I'd suggest discussing for instance 6. how the assumption that actions only have a direct effect on one object at a time limits the generality of the framework. 1. how the framework could be applied, if the sets of classes and attributes are not available

Summary

The authors has proposed Dynamic Attribute FacTored RL (DAFT-RL). For each class of object, class template graphs and interaction pattern graph and interaction graphs are computed. Through this template world model a policy can be learned for direct application in a new environment by just estimating the interactions and latent parameters. Authors showed that DAFT-RL outperforms the state-of-the-art in three compositional generalization benchmarks. Ablation studies

Strengths

The paper brings in the philosophy of class template graphs and interaction pattern graph and interaction graphs for multi-object RL. This is overall well written paper with good supplementary material details. The results validate the claims.

Weaknesses

No video demo or code is provided. The definition looks cluttered line 163. A pictorial example at the intro of a sample task is expected. (borrow from supplementary add here)

Questions

Why the same paper content is repeated in supplementary upto pg 9? remove. I was wondering why Dynamic Bayesian Networks (DBNs) was chosen? How can standard knowledge graph and graph network theory approaches be leveraged in this context?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, 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

How this method will leverage for multi object multi instance objects with ambiguity?

Summary

The paper presents DAFT-RL, which is a object-based dynamics model for reinforcement learning that learns factorized relationships between attributes of objects. These relationships are learned in several steps. First, single-object episodes are created using a random policy, and these are used to train a model of how the attributes of each object class interact with each other over time. These graphs are then held fixed, and multi-object interaction data is used to train additional graphs representing interactions between the attributes of each pair of classes, and a graph representing which object instances interact with each other. These learned graphs then help train a next-state dynamics model that predicts each object's attribute values conditioned on the relevant attribute values read off from the graphs. Finally, this model is used to either train a policy, or for a planning method like Model Predictive Control. DAFT-RL is evaluated in several experiments, where the number of objects and the latent factors of thos objects change. The paper highlights that DAFT-RL performs favorably compared to other object-centered reinforcement learning approaches on three simulated robotic benchmarks.

Strengths

- The proposed model exploits object structure that exists in many RL problems to get better performance than a broad set of baselines. To my knowledge, this is a novel approach to an interesting problem. - DAFT-RL can generalize to environments with different numbers of objects than what was seen during training. - Ablation studies are done to show relative importance of each piece of the proposed approach

Weaknesses

- Clarity of writing can be improved. In particular, Section 2 and Figure 1 were difficult to follow because the running example was not explained clearly. $C_1.s_3$ for example is not defined, and context suggests there are only two attributes of boxes. The appendix clears up the issues, so this may be fixable with some rewriting. - The multi-phase training seems a bit cumbersome. How realistic is it to be given interactions with single objects before observing multi-object interactions? - It is unclear how DAFT-RL would scale to scenes with many classes of objects, or objects with many attributes. Would the graphs become very difficult to learn? - One of the main contributions is the decomposition of object interactions to the interactions of individual attributes. Is there some ablation that could be done to analyze the benefits over modeling object-level interactions only?

Questions

- What attributes are extracted from objects in the pixel observation case? - What might the latent parameters of objects be in pixel observation environments, and would it be difficult to learn meaningful latent representations from a small set of data? - How sensitive is the RL performance to the quality of the graphs extracted in the first two phases of training? On a similar note, how sensitive are the graph learning steps to the amount of data used to train them? I imagine there may be situations where it is difficult to get a full sense of the class template graphs with data coming only from a random policy.

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

2 fair

Contribution

2 fair

Limitations

The paper would improve from a more forthright examination of the limitations of the method. How does the method scale to more complicated tasks with more objects? What is holding the method back from being able to train end to end?

Reply to the response of the author

Thank you for your response to my concerns. They properly discussed my concerns.

Slight increase

I thank the authors for their detailed feedback. I've read the rebuttal and the response to other reviewers. I think the additional points of clarification will make the paper stronger (and so I'll bump up my recommendation to weak-reject), but overall the technical contribution is marginal (and I'm not overly convinced by the rebuttal re: "extension to Object-Oriented (PO)MDPs") and I'm still unconvinced this will work on "real world" problems outside of the toy environments shown. "one could say that factored MDPs are just an extension of MDPs with factors/attributes, and hence also object-oriented MDPs are just an extension of MDPs with objects" There's false equivalency here. factored MDPs at least introduce a new concept. There's no new concept here, just an engineered extension to Object-Oriented (PO)MDPs. I don't find the core idea novel or surprising. "Second, we think our approach neatly combines several ideas from disparate communities: starting from Factored MDPs and Object-Oriented MDPs, to dynamic Neural Relational Inference for modeling interactions, soft attention networks for action binding, object-centric representation learning for extracting attributes from pixels and latent parameters that vary across objects from factored adaptation literature." Yes, I agree. It's a bunch of existing ideas put together in a not overly novel way. "We would be happy if the reviewer would point us to a real-world setting, so we can use it in our final version." This is a cop out. Look at literally any paper from CoRL 2022 (or any real-world robotics conference) if you need inspiration for what "real world" problems look like. There are many such benchmarks. You're also able to run baselines on a benchmark you propose, you're not limited to the set of environments the baselines were run on. I'm going to unfairly paraphrase the arguments above as: "look, there's precedent for other papers that only include simulated results". The problem I'm trying to convey is that your baselines didn't propose something that is a) clearly tailored to the set of simulated benchmarks you're evaluating on (with a scripted set of attribute classes that works only for those baselines) and b) is unlikely to scale to real-world problems (IMO). Therefore I'd say the burden is on this work to apply the method to a real problem and demonstrate that it outperforms the baselines on that real problem.