Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: 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.
strength_and_weaknesses: Strengths:
+ The paper is well written and easy to follow.
+ I like the theoretical framework proposed by the authors and the subsequent learning framework built on top.
+ Proposed experiments confirm the hypothesis put forward by the authors. 

Weaknesses:
- Definition 1: \psi is composed of (\psi_{int}, \psi_{ext}).
It is clear that \psi_{int} is equivariant, i.e. z'_{int} = z_{int} + a
But \psi_{ext} doesn't appear to be equivariant.
z'_{ext} != z_{ext} + a
As 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.
In 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.

- Based on previous comment, Theorem 4.1 seems to apply to \psi_{int} only and not \psi_{ext}.
This 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}.
There is no term ensuring linearity on \psi_{ext}, which is also not true in the real word.
What makes z_{ext} and z'_{ext} consistent with each other under action a.

- 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?

- According to theorem 4.1, if our model has learnt to satisfy the 3 conditions then:
\forall s_i \in D_{test} \psi(w(s_i)) - s_i = h
Why can't we evaluate how much does our model deviate from this constant h.
L_{test} = var(\psi(w(s)) - s). This value should ideally be 0. Why is current metric better?

Clarifications:
- Definition 1: Does the equivariance constrain make the latent representation z_{int} linear?
Does this affect/ limit the properties that the model can capture? Is linear representation space sufficient for complex properties?
In 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.
What if the model has to learn non-linear properties like rotations? Will it affect the quality of performance?


Minor:
- 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.
- Why consider only translation? Rigid objects can rotate too. Modelling rotations from image observations can be challenging.
- Sec 3: What is n? Is n=3? Translations in 3D Euclidean space?
- Assumption 3.1: (.) is an operator that applies actions 'a' on state 's' to give the new state s'
At the end of paragraph a is applied, using (.) operator, to the observation o. What is the domain of operator (.), state or observation?
- 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.
clarity,_quality,_novelty_and_reproducibility: The paper is well written and I enjoyed the problem setting and the accompanying formulation. The experiments support the claims made in the paper.
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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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. 
strength_and_weaknesses: # Strengths

* 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.
* I have checked the theoretical foundation and proof (to the best of my ability) and they are technically correct.
* The code has been provided for reproducibility.
* 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.

# Weaknesses

## Some assumptions are unjustified and could be limiting
* 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)?
* 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).
* 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?

## Unclear explanation of Equation 7
* 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+. 

## Evaluation metric entangles inference of z_int and z_ext
* Is it possible to independently evaluate z_int and z_ext?
* 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). 
* It is unclear how much of an error exists in inferring agent state vs. object state. 

## VAE converges to comparable RL performance while being very poor at inferring states
* 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. 
* 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.

## One point not clear in proof of theorem A1
* 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?


clarity,_quality,_novelty_and_reproducibility: The paper clarity was generally good (barring a few points raised in weaknesses). The work appears to be novel and reproducible. 
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.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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.
  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.

strength_and_weaknesses:   - Strength
    - The experiments on the Soccer environment show that the proposed method can extract plausible representations under projected observations.
    - The experiments on the control task demonstrate the quality of the inferred agent and object states.
  - Weakness
    - Interaction between objects is limited to simple collisions. Thus, it is not clear to me how general this framework is.
    - It is not clear to me the precise meaning of "geometric representation". The paper mainly focuses on the location of the agent and object.
    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.
    - The true state dimension and the form of action (additive in this case) are assumed to be known.
    - The experiments are overall weak.
      - 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.
      - For the control task, judging by the reward plot, the training does not converge yet.
clarity,_quality,_novelty_and_reproducibility:   - Not sure why the caption of figure 4 says "on both versions of the Sprites dataset".
  - Sec 5.3, why the inaccessible state space is 2D?
  - Since the paper only focus on Euclidean state space, why does figure 1 characterize the state space as a saddle surface?
  - How is the action defined? Is it simply the offset of the agent?
  - Doesn't the Soccer environment violate the assumption that the state of object remains constant when not acted upon? Why the training objective still makes sense?
  - The visualization shows that the learned visualization accurately capture the true state. How is the translation ambiguity is handled?

summary_of_the_review:   The paper in its current state demonstrates the potential of this approach but did not fully justify it.
  I think many assumption/limitation of this approach requires in-depth discussion.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: 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.
strength_and_weaknesses: Strength: 
Geometry-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.

Weaknesses:
1. 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.

2. 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.

3. 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.
clarity,_quality,_novelty_and_reproducibility: The paper is clearly written and results / derivations are well presented. The python code for reproducing the results is included in supplementary materials. To the reviewer's best knowledge, the paper sufficiently cite relevant papers in this domain, even though the reviewer is not up to date with the most recent publications.
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.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: The paper proposes a representation learning framework for the “geometric” state of the agent and objects. The proposed theory shows that an ideal learner disentangles the agent from the object in the representation space. The supervision for representation learning comes from the taken actions (the dynamics is unknown). 

Strength:
- The paper provides a theoretical formulation for representation learning (however for a very constrained setting studied in the paper).

- The paper is well-written and it is clear for a wider audience in the ML community. 

Weaknesses:
- The paper considers only the “translation” of the agent and the object ignoring all other important aspects of agent-object interaction.

- It assumes the full observability of the environment.

- It hardcodes a policy that forces the agent to interact with objects.

- The experiments have been performed in very simple environments.

justification_for_why_not_higher_score: The provided theory is only valid for a very constrained setting that involves several strong assumptions. The paper requires significant and fundamental changes to relax these assumptions.
Decision: Reject