Learning Geometric Representations of Interactive Objects
Authors: Alfredo Reichlin
Giovanni Luca Marchetti
Hang Yin
Anastasia Varava
Danica Kragic
Venue: Submitted to ICLR 2023 (2023)
Keywords: Representation Learning
Interaction
Equivariance

Abstract: We address the problem of learning geometric representations from observations perceived by an agent operating within an environment and interacting with an external object. To this end, we propose a representation learning framework that extracts the state of both the agent and the object from unstructured observations of arbitrary nature (e.g., images). Supervision comes from the performed actions alone, while the dynamics of the object is assumed to be unknown. We provide a theoretical foundation and formally prove that an ideal learner is guaranteed to infer an isometric representation, disentangling the agent from the object. Finally, we investigate empirically our framework on a variety of scenarios. Results show that our model reliably infers the correct representation and outperforms vision-based approaches such as a state-of-the-art keypoint extractor.   


Under review as a conference paper at ICLR 2023
LEARNING GEOMETRIC REPRESENTATIONS OF INTERACTIVE OBJECTS
Anonymous authors Paper under double-blind review
ABSTRACT
We address the problem of learning geometric representations from observations perceived by an agent operating within an environment and interacting with an external object. To this end, we propose representation learning framework that extracts the state of both the agent and the object from unstructured observations of arbitrary nature (e.g., images). Supervision comes from the performed actions alone, while the dynamics of the object is assumed to be unknown. We provide theoretical foundation and formally prove that ideal learner is guaranteed infer an isometric representation, disentangling the agent from the object. Finally, we investigate empirically our framework on variety of scenarios. Results show that our model reliably infers the correct representation and outperforms vision- based approaches such as state-of-the-an keypoint extractor.
INTRODUCTION
A fundamental aspect of intelligent behavior by part of an agent is building rich and structured representations of the surrounding world (Ha & Schmidhuber (2018)) Through structure, fact, a representation can potentially achieve semantic understanding, efficient reasoning and general- ization (Lake et al. (2017)). However, in realistic scenario an agent perceives unstructured and high-dimensional observations of the world (e.g.. images). The ultimate goal of inferring rep- resentation consists thus of extracting structure from such observed data (Bengio et al. (2013)) This is challenging and in some instances requires supervision or biases. For example, it is known that disentangling factors in data is mathematically impossible in a completely unsupervised way (Locatello et al. (2019)). In order extract structure, it is thus necessary to design methods and paradigms relying on additional information and specific assumptions.
In the context of an agent interacting with the world, fruitful source of information is provided by the actions performed and collected together with the observations. Based on that, several recent works have explored the role of actions in representation learning and proposed methods extract structure from interaction (Kipf et al. (219)); } Mondal et al.] (2020): Park et al. (2022). The common principle underlying this line of research encouraging the representation
replicate the effect In particular, it has been shown in Marchetti et al. (2022) that equivariance enables to extract the internal state of the agent (i.e., its pose in space), resulting in a lossless and geometric representation. The question of how to represent features of the world which are extrinsic to the agent (e.g.. objects) has been left open. Such features are dynamic since they change as consequence of interaction. They are thus challenging to capture in the representation but are essential for understanding and reasoning by part of the agent.
of actions in a structured space a property referred to as equivariance
In this work we consider the problem of learning a representation of an agent together with an exter- nal object the agent interacts with (see Figure We focus scenario where the object displaces only when it comes in contact with the agent, which realistic and practical. This enables to design a representation learner that attracts the representation of the object to the one of the agent when inter- action happens or keeps it invariant otherwise. Crucially, we make no assumption on the complexity of the interaction: the object is allowed to displace arbitrarily and its dynamics is unknown. All the
Alternative terminologies from the literature are World Model Kipf et al. 2019 and Markov Decision
Process Homomorphism (van der Pol 2020

---

Under review as a conference paper at ICLR 2023
ext
0
No interaction
Interaction
0
A
Sint
Our framework enables to learn a representation yp recovering the geometric and disen- tangled state of both an agent mm. white) and an interactable object (Zex brown) from unstructured observations o (c.g., images). The only form of supervision comes from actions a, performed by the agent, while the transition of the object (question mark) in case of interaction is unknown. In case of no interaction, the object stays invariant.
Figure I:
losses optimized by our learner rely on supervision from interaction alone i.e., on observations and performed actions. This makes the framework general and in principle applicable to observations of arbitrary nature. We moreover provide formalization of the problem and theoretical grounding for the method. Our core theoretical result guarantees that the representation inferred by an ideal learner recovers both the ground-truth states up a translation. This implies that the representation isometric (i.c., fully extracts the geometric structure of states) and disentangles the agent from the object. As consequence, the representation preserves the geometry of the state space underlying observations. The preservation of geometry makes the representation lossless, interpretable and dis- entangled. We empirically show that representations not only outperform in quality of structure a state-of-the-ee keypoint extractor, but can be leveraged by downstream learner in order to solve control tasks efficiently. summary our contributions include:
A representation learning framework extracting geometry from observations in the context an agent interacting with an external object.
A theoretical result guaranteeing that the above learning framework, when implemented by an ideal learner, infers an isometric representation for data of arbitrary nature.
An empirical investigation of the framework variety of environments with comparisons to computer vision approaches (i.e., keypoint extraction) and applications to a control task.
We provide Python code implementing our framework together with all the experiments as part of the supplementary material.
2 RELATED WORK
Equivariant Representation Learning. Several recent works have explored the idea of incorpo- rating interactions into representation learning. The common principle is infer representation which is equivariant i.e., such that transitions in observations are replicated as transitions in the latent space. One option to learn the latent transition end-to-end together with the representa- tion (Kipf et al. (2019); : van der Pol et al.] (2020); Watter et al.) (2015)). This approach is however non-interpretable and the so-obtained representations are not guaranteed to extract any structure. Alternatively, the latent transition can be designed priori. Linear and affine latent transitions have been considered Guo et al. 2019). Mondal et al. (2020) and Park et al. (2022) while transitions defined by (the multiplication of) a Lie group have been discussed in Marchetti et al. (022)) Mon- dal et al. 2022). As shown in Marchetti et al. (202)), for static scenarios (i.e., with no interactive external objects) the so-obtained representations of the underlying state of the agent. Our framework adheres to this line of research by modelling the latent transitions via the additive Lie group R.,. We however further extend the representation
structured and completely recover the geometry
2

---

(1) where Sint is the space of states of the agent (internal states) and Sext is the space of states of the object (external states). We identify both the agent and the object with their position in the ambient space, meaning that 8.m. R=" 2 8ea- The actions that the agent performs are displacements of its state i.e., the space of actions consists of translations A Re. In our formalism we thus abstract objects as material points for simplicity. The practical extension of our framework to volumetric Under review as a conference paper at ICLR 2023
include external objects. Our framework thus applies to more general scenarios and dynamies while still benefiting from the geometrical guarantees.
Keypoint Extraction. When observations are images computer vision offers a spectrum of classi- cal approaches to extract geometric structure. In particular, extracting keypoints enables to identify any object appearing in the observed images. Popular keypoint extractors include classical non- parametric methods [Lowe] (1999). Bay al. (2006)) well as modern sule--uper learning approaches (Kulkarni et al. (2019). Gopalakrishman et al. (2020). However, keypoints from an image provide representation based on the geometry of the field of view or, equivalently, of the pixel plane. This means that the intrinsic three-dimensional geometry of states of objects is not preserved since the representation differs from by an unknown projective transformation. In spe- cific situations such transformation can still be recovered by processing the extracted keypoints. This is the case when images are in first person view w.r.t. the observer: the keypoints can then be converted into three- fiee-dimmnsionall landmarks via methods such as bundle adjustment (Tiggs et al. (1999). Schonberger & Frahm (2016)).
Differently from computer vision approaches our framework is general, data-agnostic and does not rely on specific priors tied the nature of observations. Assuming that the external object in- teractable (in a particular sense), our idea is to extract geometric structure from performed actions based on functional properties alone (see Section 4.1).
Active Perception. The role of interaction in perception has been extensively studied in cognitive
sciences and neuroscience (Held & Hein [963); Gibson & Carmichael| (1966); Noë et al.) (2004) )
Inspired by those, the field of active perception from robotics aims enhance the understanding the world by part of an artificial systems via interaction (Bohg al. ((017)) Applications include active control of cameras (Bajesy) (1988) and manipulators (Tsikos & Bajesy (991]) order to improve the perception of objects (llonen et al. (2014); Björkman et al. (2013); Schiebener al. (201)) Our work fits into the program of active perception since we rely on interaction in order to infer the representation. We show that actions enable to extract geometric and disentangled states from general and unstructured observations, which would not be possible without the the information gained from interaction.
FORMALISM AND ASSUMPTIONS
In this section we introduce the relevant mathematical formalism together with the assumptions necessary for our framework. We consider the following scenario: agent navigates in a Euclidean space and interacts in an unknown way with an external object. This means that the space of states 8 is decomposed as
8 = Sint x Sext
objects is discussed in Section 4.3 while the extension of agent's 's actions to arbitrary Lie groups is discussed Section
Our first assumption (comnectedness) is that the agent can reach any position from any other via sequence of actions. This translates in the following topological condition:
Assumption 3.1. (Connectedness) The space Sint connected and open.
When the agent performs an action A the state (sime Sex) transitions into state denoted by a = (s' mp s'xx) Since the actions displace the agent, the internal state gets translated as s = int + a
However, the law governing the transition of the object S'x, = T(s, a) is assumed to be unknown and can be arbitrarily complex and stochastic. We stick to deterministic transitions for simplicity of explanation. Crucially, the agent does not have access to the ground-truth state
Whenever we write implicitly assume that the action is valid i.c., that sm + E Sma-
3

---

Under review as a conference paper at ICLR 2023
Instead it perceives unstructured and potentially dighedimensional observations images) via an unknown emission map w - O. We assume that w is injectivite so that actions induce deterministic transitions of observations. The latter denoted by abuse of notation as = a o. This assumption is equivalent to total observability of the scenario and again simplifies the forthcoming discussions by avoiding the the need to model stochasticity.
(e.g.,
The fundamental assumption of this work that the dynamics of the external object revolves around contact i.e., the object does not displace unless it is touched by the agent. This natural and often satisfied in practice. In order to formalize it, note that when the agent in state ii performs an action we can imagine it moving along the open segment Sint iin + a assumption then translates into (see Figure for a graphical depiction):
(s + to Ja]uu<<1 Our
Assumption 3.2. (Interaction Occurs at Contact) For all agent states 8im E S and actions E holds that s'xt eex if and only if Sex |
int + a.-
As such, the dynamics of the external object can be summarized as follows:
Sext (s)))) otherwise.
if Sext & [smm Sint + aal,
s'cu =
(2)
Finally, we need assume that interaction is possible for every state of the object i.e., the latter has to be always reachable by the agent. This is formalized via the following inclusion: Assumption 3.3. (Reachability) holds that Sent
Sint-
4 METHOD
4.1 REPRESENTATIONS AND EQUIVARIANCE
We now outline the inference problem addressed in the present work. Given the setting introduced in Section 3 the overall goal is to infer a representation of observations
:: 0 2 Zim X Zext
(3)
where Zmm Zext R". Ideally & recovers the underlying inaccessible state in 8 C 2 and disentangles Sin from Seat- In order to achieve this, our overall idea is to split the problem of representing the agent and the object. Since the actions of the agent available Fint can be inferred geometrically by existing methods from the literature. The representation of the object, =ex E Zext. can then be inferred based on the one of the agent by exploiting the relation between the dynamics of the two (Equation 2). In order to represent the agent, we consider the fundamental concept of (translational) equivariance:
Definition 1. The representation yp is said to be equivariant (on internal states) if for all E A and o E O it holds that 2/m iint + a where (Zmm Fexa) = ((() and (z/m =em) = p(aa. o).
We remark that Definition I Tefeers to internal states only, making our terminology around equivari- ance unconventional. As observed in previous work (Marchetti et al. (2022). equivariance guaran- tees a faithful representation of internal states. Indeed if is equivariant then žint differs from iim by a constant vector. This means that the representation of internal states a translation of ground-truth ones and as such is lossless (i.e., bijective) and isometrically recovers the geometry of Sint-
The above principle can leveraged in order to learn representation of external states with the same benefits as the representation of internal ones. The intuition is that since the external object displaces only when it comes contact with the agent (Assumption (.2),, Zext can be inferred by making it coincide with int The following theoretical result formalizes the possibility of learning such representations and traces the foundation of our representation learning framework.
Theorem 4.1. Suppose that the representation &p O - Z satisfies:
1. &p is equivariant (Definition) 7).
2. yp is injective,
3. for all o E O and a E A it holds that z'ou # Zea if and only if eea E L 2 + a] where (Z Foxt) (p()) and (2/ms z(x)) = sp(a 0).
4

---

Under review as a conference paper at ICLR 2023
Then yow is translation i.e., there is constant vector h R" such that for all s E 8 it holds that ((w(ss) = s + h. In particular, yp is an isometry w.t.. the Euclidean metric on both 8 and 2.
We refer to the Appendix for proof. Theorem states that if the conditions 1. - 3. are satisfied (together with the assumptions stated in Section 3 then the representation recovers the inaccessible state up to translation and thus faithfully preserves the geometry of the environment. All the conditions from refer to properties of depending on observations and the effect of actions on them, which are accessible and collectable in practice. The corresponding properties can be enforced to y by optimizing a loss, which will be discussed in the next section.
4.2 LEARNING THE REPRESENTATION
In this section we propose representation learning framework implementing the conditions from Theorem 4.1 via number of losses. We assume that the representation y = (fiitty (exx) consist of two parameterized functions (Pm
-> Zins ffex : -> Zeu e.g., two deep neural network models. In order to design the losses that the model(s) minimize w.r.t the respective parameters, we consider a dataset of transitions D = {(,, a, of = a - }}} observed by the agent. Such is collected by the agent exploring the environment and interacting with the external object described in Section 3 We remark that our representation learning framework relies on supervision from interaction alone, without access to the ground-truth states of the external object nor of the agent itself.
First, we propose to enforce equivariance, condition 1 from Theorem by minimizing the loss:
Lim (o, a, o') = d(2/mt Pint - a)
(4)
where d is a measure of similarity on Zjnt = Rh and the notation is in accordance with Definition [ Typically d is chosen as the squared Euclidean distance. Equation canonical loss which
has been considered in previous work addressing equivariant representation learning (Mondal et al. 2020 Kipf et al. (2019))
Next, we focus on condition 3. The idea is to force (Pet to minimize one of the following losses depending on the given datapoint:
f- o') = d( d(=exts =(o)))
ao,00,,'') d(zem L iint + aa)-
(5)
The distance involved in L++ represents point-to-set metric and typically set as d(z, F)) = inf e d(z, x). The latter has a simple explicit expression in the case E is segment. One nat- ural way to implement condition 3 is to combine the two losses from Equation 5 into a single one as:
L Lo a, o') = minf (- (lo,,,,')) o'), Lo a, o')
(6)
However, we have found that co is hard to optimize in practice. This happens because & forces p to choose between &_ and 64+ and, at the same time, to optimize the chosen one, resulting in an complex task. Instead, we rely the following way of combining the two losses which leads to an easier optimization problem. We propose to deploy a latent contrastive representation in order decide which one among L_ and L4 minimize for a given datapoint. Namely we train
another model Feom O + W with latent space W (potentially different from 2) which attracts
Peout (o) to w' = Peon((o') Given distance dw. we stick to the popular InfoNCE loss from contrastive learning literature (Chen et al. (2200)):
Lo o') dw w') + log Eg" [e (ww 20") ) --d(= "
where o" is marginalized from D. The second summand of Equation 7 encourages the joint encod- ings (zmm w) to spread apart. The latter avoids the collapse of PPoon to a constant and encourages condition from Theorem 4.1 (injectivityy) A typical choice to normalize the output of Poont so that W is sphere and set dw (w, w') = cos(Lww') = -w-wu ww' (Wang & Isola (2020)). The latent space W enables to choose whether to optimize C_ or L since subsequent observations where the agent does not touch the object will lay close in W. We propose to partition (the given batch of) the dataset in two disjoint classes D = C_ Li C+ by applying a natural thresholding algorithm to the quantities dw (w, w²". This can be achieved via one-dimensional 2-means clustering, which is equivalent to Otsu's 's algorithm (Otsul (179)). We then optimize:
a (c_(o,,,o') if (o, CC,, (0,00,0,0) if(o,a,)) E C+.
(8)
5

---

Under review as a conference paper at ICLR 2023
In summary, the total loss minimized by the models (pmms (Pex, (Poou) w.r.t. the respective parameters is (see the pseudocode included in the Appendix):
L -(o,a,0)) a, d') + Lo, a, d') (com(o,
4.3 INCORPORATING SHAPES OF OBJECTS
So far we have considered the external object as a point in Euclidean space. However, the object typically manifests with body and thus occupies volume. Interaction and consequent displace- ment (Assumption 3.3) occurs when the agent comes in contact with the boundary of such volume. In order represent states faithfully, the representation has take the volume of the object into account. We propose to naturally achieve this via stochastic outputs i.e., by designing Zext as probability density over Zex representing (a fuzzy approximation of) the body of the object. More concretely, the output of fPet consists of (parameters of) Gaussian distribution whose covariance matrix represents the body of the object. The distance appearing in Equation 5 is replaced with Kullback-Leible divergence. The latter has explicit simple expression for Gaussian densities which allows to compute L_ directly. In order to compute 4+ we rely on a Monte Carlo approx- imation, meaning that sample point uniformly from the interval and set L as its negative log-likeli log-likelhhod w.r.t. the density of žext
EXPERIMENTS
We empirically investigate the performance of our framework in correctly identifying the position of agent and of an interactive object. The overall goal of the experimental evaluation showing that our representation is capable of extracting the geometry of states without relying on any prior knowledge of observations e.g., depth information. All the scenarios are normalized so that states lie in the unit cube. Observations are RGB images of resolution 100 x 100 in all the cases considered. We implement each of (Pint, Pext and &Poont as ResNet-18 (He et al. (2016) and train them for 100 epochs via the Adam optimizer with learning rate 0.001 and batch-size 128.
5.1 BASELINES AND EVALUATION METRICS
We compare our framework with two baselines:
Transporter Network (Kulkarni et al. (019)): ): a vision-based state-of-the-art unsupervised keypoint extractor. The approach heavily relies on image manipulation in order to infer regions of the pixel plane that are persistent between pairs of images. We train the model in order to extract two (normalized) keypoints representing Ann and Zext respectively.
Variational AutoEncoder (VAE) (Kingma & Welling (203)); Rezende et al. (2014)); a pop- ular representation learner with standard Gaussian prior on its latent space. We impose the prior on Zes only, while Pmm is still trained via the equivariance loss (Equation E The decoder takes the joint latent space Z in input. We set dim(Zex) = 32. This makes the representations disentangled, so that im and z well-defined. The resulting represen- tation of the object generic and is not designed to extract any specific structure from observations.
In order to evaluate the preservation of geometry rely the following evaluation metric Liest- Given a trained representation O - and a test set Dress of observations with known ground- truth states, we define:
Lest = Bog-Dig [(( Fextr Sint *
(10)
where d is the squared Euclidean distance. Since both our framework and (the encoder of) VAE have stochastic outputs (see Section 4.3), we set Zeu the mean of the corresponding Gaussian distribution. Equation 10 measures the quality of preservation of the relative position between the agent and the object by part of the representation. When Leest 0, y is an isometry (w.Lt. the Euclidean metric) and thus recovers the geometry of states. The translational invariance of Liest makes the comparison agnostic to any reference frame eventually inferred by the given learner.

---

Under review as a conference paper at ICLR 2023
a
Trinpprter
Figure 2: Top: Visualization of the dataset from the Sprites experiment. On the left, an example of observation o, action and the resulting observation o'. On the right, an example of an observation from the second version of the dataset where a moving background added as a visual distractor. Bottom: Comparison of Zim Fext (gray dots, with the ellipse representing the learned std) extracted via our model and the Transporter network on both versions of the Sprites dataset.
5.2 SPRITES
For the first experiment we procedurally generate images two sprites (the agent and the object) moving black background (see Figure 2). Between images, the agent (red figure) moves cording to a known action. If the agent comes in contact with the object (green diamond) during the execution of the action (see Assumption 3.2) the object is randomly displaced the image. In other words, the object's transition function l's a) is stochastic with a uniform distribution. Such completely stochastic dynamics highlights the independence of the displacement of the agent w.r.t. the one of the object. We additionally consider a second version of the experiment by simulating moving background. Images are now overlayed on top three-times larger second image (blue squares in Figure [ right). The field of view and thus the background moves together with the agent. The moving background behaves as visual distractor and makes challenging extract structure (e.g., keypoints) via computer vision.
Transporter
VAE
Log Sin
lug Lus
100
Epochs
Epochs
Figure 3: Log-scale plots of the evaluation metric (Equation 10 as the training progresses for the Sprite experiment. The curves display mean and std (for 10 experimental runs). Left: standard version of the dataset. Right: challenging version with moving background.

---

Under review as a conference paper at ICLR 2023
Figure 3 displays the analytic comparison of the performances between our model and the baselines in terms of the evaluation metric (Equation [Op. The plot is in log-scale for visualization purposes. Moreover, Figure 2 (bottom) reports qualitative comparison between model and the Transporter network. As can be seen, for the simpler version of the experiment (plot on the left) both our model and the Transporter network successfully achieve low error and recover the geometry of both the agent and the object. Note that the Transporter network converges slowly and with high variance (Figure 3 left). This is probably due to the presence of a decoder in its architecture. Our framework instead involves losses designed directly in the latent space, avoiding an additional model to decode observations. As expected, VAE achieves significantly worse performances because of the lack of structure in its representation.
For the more challenging version of the experiment, the transporter not able to extract the expected keypoints. The moving background causes the model to focus on regions of the image not corre- sponding to the agent and the object. This is reflected by significantly higher error (and variance) w.r.t our framework (Figure 3 right). The latter still infers the correct representation and preserves geometry. This empirically confirms that our model is robust to visual distractors since it does not rely on any data-specific feature or structure.
5.3 SOCCER
For the second experiment we test our framework on an environment consisting of an agent on a soccer field colliding with ball (see Figure E. top-left). The scene is generated and rendered via the Unity engine. The physies of the ball is simulated realistically: in case of contact, rolling takes gravity and friction into account. Observations are third-person views of the scene from a fixed camera. Note that even though the scene is generated via three-dimensional rendering, the (inaccessible) state space is still two-dimensional since the agent navigates on the field. Our dataset consists of 10000 generated triples (o, a, o' = a a....
Figure (top-right) displays the comparison of the performances between our model and the base- lines in terms of the evaluation metric (Equation I0). On the bottom-left, we report visualizations of the learned representations. The Transporter network correctly extracts the keypoints on the pixel plane. As discussed in Section 2 such plane differs from Sin by an unknown projective (and thus non-isometric) transformation. This means that despite the successful keypoint extraction, the ge- ometry of the state space not preserved, which is reflected by the high error on the plot. This is a general limitation of vision-based approaches: they are unable recover the intrinsic geometry due to perspective in case of a three-dimensional scene. Differently from that, our framework correctly represents the scene and achieves low error independently from observations.
5.4 CONTROL TASK
In our last experiment we showcase the benefits of our representations in solving downstream control tasks. The motivation is that a geometric and dineennnnional representation improves efficiency and generalization compared to solving the task directly from observations. To this end we design control task for the Soccer environment consisting in kicking the ball into the goal. The reward is given by the negative distance between the (barycenter of the) the ball and the (barycenter of the) goal. In each episode the agent and the ball is initially placed in random location while the ball is placed in the center. The maximum episode length is 20 steps. We then train a number models via the popular reinforcement learning method Proximal Policy Optimization (PPO; (Schulmculmaa et al. 2017)) ). One model (End-to-End) receives raw observations as inputs. The others operate on pre- trained representations Z given by the Transporter network, the VAE and our method respectively. All the models implement comparable architecture for a fair comparison.
Figure 4 (bottom-right) displays the reward gained on test episodic runs as the training by reinforce- ment learning progresses. As can be seen, our geometric representation enables to solve the task more efficiently than both the competing representations (Transporter and VAE) and the end-to-end model. Note that the Transporter not only does not preserve the geometry of the state space, but has the additional disadvantage that the keypoint corresponding to the agent and the object can get swapped in the output of yo This causes indeterminacy in the representation and has negative im- pact on solving the task. Due to this, the Transporter performs similarly to the end-to-end model and outperformed by the generic and non-geometric representation given by the VAF. In conclusion,
8

---

Under review as a conference paper at ICLR 2023
Ours
Transporter
End-to-EEd
log, Eres
Epochs
Ours
Reward
Trmsporter
20:00
500ce
43000
55300
Steps
on the left, an example of an observation from the Soccer experiment. On the right, log-scale plot of the evaluation metric the training progresses. Bottom: the left, visual comparison of sm eext (red dots) extracted via our model and the Transporter network. For our model, visualize view of the scene from the top instead of the original observation. On the right, plot of the reward gained via reinforcement learning on top of different representations.
Figure 4: Top:
the results show that downstream learner can significantly benefit from geometric representations of observations in order to solve downstream control tasks.
CONCLUSIONS AND FUTURE WORK
In this work proposed novel framework for learning representations of both agent and an object the agent interacts with. We designed a system of losses based on a theoretical principle that guarantees isometric representations independently from the nature of observations and relying supervision from performed actions alone. We empirically investigated our framework on multiple scenarios showcasing advantages over computer vision approaches.
Throughout the work we assumed that the agent interacts with single object. An interesting line of future investigation is extending the framework take multiple objects into account. In the stochastic context (see Section 4.3) an option is to model eau via multi-modal densities, with each mode corresponding to an object. This however comes with a computational challenge: since the dataset has to cover transitions densely for a successful inference, the sample complexity would grow exponentially w.r.t. the number objects. order to design sample-efficient method. Similarly, another possible extension of our model is considering multiple agents simultancously interacting with external objects.
improvement would thus be necessary in
As additional line for future investigation, our framework can be extended to actions beyond translations Euclidean space. Lie groups other than R" often arise in practice. For example, the agent is able to rotate its body then (a factor of) the space of actions has to contain the group of rotations SO(n). n = 2, Thus, a framework where actions (and consequently states) are represented in general Lie groups defines useful and interesting extension.

---

Under review as a conference paper at ICLR 2023
7 REPRODUCIBILITY STATEMENT
We include Python code for all the experiments and baselines the supplementary material. For the Sprites experiment, we include the code for generating data procedurally. For the Soccer experiment, we include an anonymized link to download the dataset and code to run the reinforcement learning models in Unity. This makes our experiments fully reproducible. For the main theoretical result (Theorem 4.1. the complete proof is included in the Appendix while all the assumptions are stated in Section
REFERENCES
Ruzena Bajesy, Active perception. Proceedings of the IEEE, 76(8):6661005, 1988.
Herbert Bay, Tinne Tuytelaars, and Luc Van Gool. Surf: Speeded up robust features. European conference on computer vision, pp. 404-417. Springer, 2006.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 35(8):1798-1828, 2013.
Marten Björkman, Yasemin Bekiroglu, Virgile Högman, and Danica Kragic. Enhancing visual perception of shape through tactile glances. In 2013 IEEE/RSI International Conference on In- doi: 10.1 SITIIROS..01..6668800
telligent Robots and Systems, pp. 3180-3186, 2013.
Jeannette Bohg, Karol Hausman, Bharath Sankaran, Oliver Brock, Danica Kragic, Stefan Schaal, and Gaurav S Sukhatme. Interactive perception: Leveraging action in perception and perception in action. IEEE Transactions on Robotics, 137)):127-1 2017.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pp. 1597-1607. PMLR, 2020.
James Jerome Gibson and Leonard Carmichael. The senses considered as perceptual systems, vol- ume 2. Houghton Mifflin Boston, 1966.
Anand Gopalakrishnan, Sjoerd van Steenkiste, and Jürgen Schmidhuber. Unsupervised object key- point learning using local spatial predictability. ICLR 2021, 2020.
Xifeng Guo, En Zhu, Xinwang Liu, and Jianping Yin. Affine equivariant autoencoder. In I/CAI, pp- 2413-2419, 2019.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Richard Held and Alan Hein. Movementtproodueed stimulation in the development of visually guided behavior. Journal of comparative and physiological psychology. 56(5):872, 1963.
Jarmo Honen, Jeannette Bohg, and Ville Kyrki. three-dimensional object reconstruction of symmet- ric objects by fusing visual and tactile sensing. The International Journal of Robotics Research, 33(2):321-341, 2014.
Diederik P Kingma and Max Welling. 1eXi.61312.611 2013.
Auto-encoding variational bayes. arXiv preprint
Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. arXiv preprint arXiv:1911.122 2019.
Tejas D Kulkarni, Ankush Gupta, Catalin lonescu. Sebastian Borgeaud, Malcolm Reynolds, Andrew Zisserman, and Volodymyr Mnih. Unsupervised learning of object keypoints for perception and control. Advances in neural information processing systems, 32, 2019.
10

---

Under review as a conference paper at ICLR 2023
Brenden M Lake, Tomer D Ullman, Joshua Tenenbaum, and Samuel Gershman. Building machines that learn and think like people. Behavioral and brain sciences, 40, 2017.
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In international conference on machine learning, pp. 41 4124-4124. PMLR, 2019.
David G Lowe. Object recognition from local inverint features. In Proceedings of the seventh IEEE international conference computer vision, volume 2, pp. 1150-1157, leee, 1999.
Giovanni Luca Marchetti, Gustaf Tegnér, Anastasiia Varava, and Danica Kragic. Equivariant repre- sentation learning via class-pose decomposition. arXiv preprint arXiv:2207.03/16, 2022.
Arnab Kumar Mondal, Pratheeksha Nair, and Kalcem Siddiqi. Group equivariant deep reinforce- ment learning. arXiv preprint arXiv:2007. 0034377 2020.
Arnab Kumar Mondal, Vineet Jain, Kaleem Siddiqi, and Siamak Ravanbakhsh. Eqr: Equivariant representations for data-efficient reinforcement learning. In International Conference on Machine Learning, pp. 15908-15926. PMIR, 2022.
Alva Noë, Alva Noë, et al. Action in perception. MIT press, 2004.
Nobuyuki Otsu. A threshold selection method from gray-level histograms. IEEE transactions on systems, man, and cybernetics, 9(1):62-66, 1979.
Jung Yeon Park, Ondrej Biza, Linfeng Zhao, Jan Willem van de Meent, and Robin Walters. Learning symmetric embeddings for equivariant world models. arXiv preprint urXiv:2204.1137,, 2022.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and ap- proximate inference in deep generative models. In International conference on machine learning, pp. 1278-1286 PMLR, 2014.
David Schiebener, Jun Morimoto, Tamim Asfour, and Aleš Ude. Integrating visual perception and manipulation for autonomous learning of object representations. Adaptive Behavior, 21(5); 328-345. 2013. doi: 10.1177/10991223 3484500. URL lattp:://aci.org///
Johannes L Schonberger and Jan-Michael Frahm. Stracture-from-mottin revisited. In Proceedings of the IEEE conference on computer vision and puttern recognition, pp. 4104-4113, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 2017.
Bill Triggs, Philip F McLauchlan, Richard Hartley, and Andrew W Fitzgibbon. Bundle ad- justment- a modern synthesis. In International workshop vision algorithms, pp. 298-372. Springer, 1999.
C.J. Tsikos and R.K. Bajesy. Segmentation via manipulation. IEEE Transactions on Robotics and Automation, 703):306-319, 1991. doi: 10.1109/70.8810.
Elise van der Pol, Thomas Kipf, Frans A Oliehoek, and Max Welling. Plannable approximations mdp homomorphisms: Equivariance under actions. arXiv preprint arXiv:2002. 1963, 2020.
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through align- ment and uniformity on the hypersphere. In International Conference on Machine Learning, pp- 9929-9939. PMLR, 2020.
Manuel Watter, Jost Springenberg. Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances neural information processing systems, 28, 2015.
11

---

Under review as a conference paper at ICLR 2023
A APPENDIX
A.1 PROOFS OF THEORETICAL RESULTS
Theorem A.1. Suppose that the representation yp: O -+ Z satisfies:
1. ap is equivariant (Definition 7
2. yp is injective,
3. for all O and A it holds that zoo # Zoa if and only if 2ext L Pint al where (a Fex)) :p() and (z/mx zex)) = ppa o).
Then yow is a translation i.e,, there is a constant vector E R such that for all s E it holds that (s((ss) + h. In particular, yp isometry w.rl. the Euclidean metric both S and 2.
Rev:
Rest
+
Figure 5: Graphical depiction of the proof of Theorem 4.1
Proof. Pick an arbitrary state o E 8 together with its represented internal state 2 and set = 20 Spa For any state s, consider the action Sint s Equivariance then implies that Smm + h. This shows that the claim holds for internal states.
Amm
am +
To prove that the same happens for external states, suppose by contradiction that there state such that Fex * Sext + h. Consider any path during which the agent interacts with the object without passing through Zext Formally, this means considering sequence of actions a1," ; ar such that (see Figure 5 left):
Zext and Sext h do not belong L + a1 + + 0i- Pint ai + as] for every
Sex does not belong to Ziint + + ... 0r Pine + an + +... a, but Sea + h does.
The existence of such path follows from Assumptions 3.1 and 3.3 After interaction the state
becomes s = a, (ar-1 ... (a, s)) with s'ca * Sext because of Assumption 3.2 One can then consider a path back the initial agent's 's position in i.c., another sequence of actions a,+1, aR such that (see Figure right):
det and
z do not belong Lzm + aj + + ai Fmm ai + + 0s. for every R.
= r + 1,
+ + an = O.
All the conditions imply together that the representation of the object remains equal zex during the execution of the actions Q1. ag. Since the actions sum to 0, the representation of the agent does change as well. But then (s)(ss) = sp(w(sin s((x)) while 8ext * s contraddicting injectivity.
We conclude that žext = ext and thus + h as desired.
12

---

Under review as a conference paper at ICLR 2023
A.2 PSEUDOCODE FOR LOSS COMPUTATION
Algorithm Loss Computation
Input: Batch B D, models Puatt Peet, Peant
Output: Loss &
== for all (o, ao) E B do
Compute Zim. = 9imi (o), Fea (ox(o), 2/m (mm((o), Lexr Pex(('), ) Peum (o), w/ = &
end for
Compute the classes C-, C+ via Otsu's algorithm based {dw (w,, w')}
for all (o, a, o') E 'B do
Compute Limlo, a, o') via Equation 4
Compute A= fdw(w', w"), d( d(zmm ( for o" marginalized from B
Based compute Looul0, o') Equation 7
if dw w. w') E C- then
Compute Lailo, 0') Lo, via Equation (left)
else
Compute Len(o, 0') L+(o, via Equation (right)
end if
1 - L + Lise + Lext + hcom
end for
13


Reviews:
[{"review_id": "fJ1nq9mKXM", "reviewer": "Reviewer_a3YM", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a framework for learning geometric-aware disentangled representations via agent-environment interaction. This goal is achieved by introducing an equivariant loss term on top of the ordinary distance losses to enforce geometric disentanglement between agent and environment state, together with a contrastive loss to assist learning. The paper provides detailed mathematical formulation of the general setup and conducts empirical studies on simple datasets whose state space are Euclidean. Experiments show preliminary comparisons between the proposed method and two other self-supervised representation learning baselines.", "summary_of_the_review": "Even though this paper conducts a close-loop study of a novel representation learning framework, I find it lacking justifications to its limitations as well as empirical studies. A lot of work is yet to be done for this paper to reach its true potential. My current assessment is reject.", "strength_and_weaknesses": "Strength: \nGeometry-aware representations are of great importance in modeling agent-environment interaction, with rich potential applications. This paper focuses on an interesting scenario of learning Euclidean positions of point-like agent and environmental objects. Experiments successfully demonstrates the geometry awareness of the learned representation, as well as its application in simple control tasks.\n\nWeaknesses:\n1. The proposed formulation seems to require the agent state and environment state to be embedded in the same geometric space. It is unclear how the method generalizes to states without a shared geometric grounding between them, such as object properties.\n\n2. The logic of the paper is not very clearly laid out, especially the relation between the theoretical derivations and empirical studies. The theory part does not take the discussion much further than the introduction of the objective functions. The paper could have been a lot stronger if more contents are dedicated to empirical studies.\n\n3. Experimental setup are too simplified (partly due to the fundamental limitation of the formulation). The paper makes claim about generality of its proposal hence it is expected to see the method plays out under more realistic environments and weaker assumptions (ideally more complex than point-objects). Two simplified environments x two baselines seem insufficient for this purpose.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "erLyUphoMH", "reviewer": "Reviewer_tBxf", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "  This paper follow the equivariant representation learning literature and proposed a equivariant representation learning framework for one-agent-one-object environment.\n  By assuming that action additively change the state of the agent and the state of the object can only be changed by agent upon impact, the proposed model managed to recovery the underlining state space.\n", "summary_of_the_review": "  The paper in its current state demonstrates the potential of this approach but did not fully justify it.\n  I think many assumption/limitation of this approach requires in-depth discussion.", "strength_and_weaknesses": "  - Strength\n    - The experiments on the Soccer environment show that the proposed method can extract plausible representations under projected observations.\n    - The experiments on the control task demonstrate the quality of the inferred agent and object states.\n  - Weakness\n    - Interaction between objects is limited to simple collisions. Thus, it is not clear to me how general this framework is.\n    - It is not clear to me the precise meaning of \"geometric representation\". The paper mainly focuses on the location of the agent and object.\n    In this case, both the agent and object can be abstracted as a point (with no meaningful rotation). I didn't find any discussion on agent or object shapes.\n    - The true state dimension and the form of action (additive in this case) are assumed to be known.\n    - The experiments are overall weak.\n      - Transporter Network by design does not consider any action, thus, it cannot filter out distractors. But I do believe if the author could visualize all key points, the agent and the object should be well captured. A fair comparison would be to apply post-processing to the set of key points and pick out the top two that are most correlated with actions.\n      - For the control task, judging by the reward plot, the training does not converge yet.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "al0VhU6jq00", "reviewer": "Reviewer_M4Jm", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper addresses the task of estimating the states (translation) of the agent and the object that it interacts with using image evidence. Authors assume that the states and properties of the agent and the object are unknown but the actions are known. The key contributions of the work is the formulation that allows decoupling of agent and object features in the latent space.", "summary_of_the_review": "I'm not an expert in this domain, so I didn't fully follow why is \\psi equivariant and not just \\psi_{int}. This is important as Theorem 4.1 (the main contribution) requires \\psi to be equivariant. See weakness section for details. Apart from this I like the work.\n", "strength_and_weaknesses": "Strengths:\n+ The paper is well written and easy to follow.\n+ I like the theoretical framework proposed by the authors and the subsequent learning framework built on top.\n+ Proposed experiments confirm the hypothesis put forward by the authors. \n\nWeaknesses:\n- Definition 1: \\psi is composed of (\\psi_{int}, \\psi_{ext}).\nIt is clear that \\psi_{int} is equivariant, i.e. z'_{int} = z_{int} + a\nBut \\psi_{ext} doesn't appear to be equivariant.\nz'_{ext} != z_{ext} + a\nAs the state of the object after the contact would depend on the shapes of the object and angle of incidence. How is \\psi_{ext}) and equivariant function.\nIn practice authors move the object at a random location, for the first experiment. Specific to this experiment: Doesn't moving the object to a random location greatly simplify the task, as now we just need to disentangle agent (whose movement follows action) and object (which is moved randomly)? More challenging is the real setting when both agent and object move according to the applied action.\n\n- Based on previous comment, Theorem 4.1 seems to apply to \\psi_{int} only and not \\psi_{ext}.\nThis can also be seen in the loss terms. L_{int} only makes \\psi_{int} equivariant and L_{ext}+L_{cont} just enforce condition 3 on \\psi_{ext}.\nThere is no term ensuring linearity on \\psi_{ext}, which is also not true in the real word.\nWhat makes z_{ext} and z'_{ext} consistent with each other under action a.\n\n- The idea behind eq. 7,8 is to train a proxy encoder \\psi_{cont} using contrastive learning and use it's prediction to generate pseudo labels, by partitioning the dataset, to optimise eq. 6. Authors do so because directly optimising eq. 6, collapses the optimisation to one of the terms (L-, L+). Why can't we enforce InfoNCE directly on z_{int} to regularise the space and spread samples in the latent space?\n\n- According to theorem 4.1, if our model has learnt to satisfy the 3 conditions then:\n\\forall s_i \\in D_{test} \\psi(w(s_i)) - s_i = h\nWhy can't we evaluate how much does our model deviate from this constant h.\nL_{test} = var(\\psi(w(s)) - s). This value should ideally be 0. Why is current metric better?\n\nClarifications:\n- Definition 1: Does the equivariance constrain make the latent representation z_{int} linear?\nDoes this affect/ limit the properties that the model can capture? Is linear representation space sufficient for complex properties?\nIn the current formulation the model is reasoning about the translation of the agent and the object, which is linear. Thus enforcing linearity on the latent space makes sense.\nWhat if the model has to learn non-linear properties like rotations? Will it affect the quality of performance?\n\n\nMinor:\n- Caption in Fig. 1: Text uses 's' to describe states of the agent and the object, but the figure caption uses 'z'. 'z' is later used as representation of observation and not the state itself.\n- Why consider only translation? Rigid objects can rotate too. Modelling rotations from image observations can be challenging.\n- Sec 3: What is n? Is n=3? Translations in 3D Euclidean space?\n- Assumption 3.1: (.) is an operator that applies actions 'a' on state 's' to give the new state s'\nAt the end of paragraph a is applied, using (.) operator, to the observation o. What is the domain of operator (.), state or observation?\n- Maybe out of scope for this work but it should be mentioned that state change can happen with a previous contact/interaction. Eg: if you push an object (like a ball) it can continue to move or stop when the contact is removed.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "2z7ypuYpOfq", "reviewer": "Reviewer_ySVW", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a representation learning framework for identifying agent and object representations from observations in the scenario of an agent interacting with an object. In particular, the work aims to learn an isometric representation (i.e., capture underlying geometric states in a loss-less fashion) with object and agent states disentangled. The work also lays the theoretical foundation for this framework and formally proves that an ideal learner can successfully learn such representations. Experiments are presented on grid-worlds and a soccer environment. ", "summary_of_the_review": "The idea is interesting, well-motivated and clearly described. The experiments are well designed and demonstrate the fundamental points shown in the theoretical framework. I find that some assumptions may be strong and limiting, and a few other weaknesses in terms of clarity of Eqn. 7, the evaluation metric, and inconsistency b/w VAE representation performance vs. RL performance. I would appreciate it if the authors could address these concerns.", "strength_and_weaknesses": "# Strengths\n\n* The idea is well motivated and clearly described. In particular, Sections 3 and 4 lay the foundation for inexperienced readers and have good clarity. The assumptions are also clearly stated.\n* I have checked the theoretical foundation and proof (to the best of my ability) and they are technically correct.\n* The code has been provided for reproducibility.\n* The experiments, while toyish, are well designed and show consistent improvements over baselines. The qualitative analysis from Figures 2 and 4 are helpful to understand the model in action.\n\n# Weaknesses\n\n## Some assumptions are unjustified and could be limiting\n* Page 4, line 3: \"We assume injectivity in the emission …\" --- doesn't this exclude any form of partial observability (for example, the object may not be visible, or the agent could be facing a blank wall in with limited field of view)?\n* Theorem 4.1, condition 3 - the \"if and only if\" excludes cases where the agent chooses not to interact with the object (i.e., the agent moves past the object without affecting it).\n* Sec. 4.2, lines 5-6: \"Such D is collected by the agent exploring the environment …\" --- this assumes that there are sufficient examples of the agent interacting with the object. How is this guaranteed while generating the training triplets?\n\n## Unclear explanation of Equation 7\n* I was not able to fully grasp how L_cont is derived in Eqn. 7. It is also unclear what the latent-space W is capturing, and how this helps select between L- and L+. \n\n## Evaluation metric entangles inference of z_int and z_ext\n* Is it possible to independently evaluate z_int and z_ext?\n* For example, we could compare (z_int@t - z_int@0) with (s_int@t - s_int@0), and (z_ext@t - z_ext@0) with (s_ext@t - s_ext@0). \n* It is unclear how much of an error exists in inferring agent state vs. object state. \n\n## VAE converges to comparable RL performance while being very poor at inferring states\n* In Fig 4  - VAE performs very poorly when it comes to inferring the agent and object states, but it eventually achieves comparable RL task reward when compared to the proposed method. \n* Does this imply that we may not need isometric representations that are translationally shifted from the agent / object states? Perhaps less structured representations are sufficient.\n\n## One point not clear in proof of theorem A1\n* Page 12 - \"All the conditions imply ... remains equal to z_ext during the execution of actions a1 ... aR\" -- why is this the case? After interaction, the object representation could have changed to some z_ext^’ != s_ext^’ + h right?\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject