Successor Representations Enable Emergent Compositional Instruction Following

Thank you for your detailed response and thoughtful feedback. We would like to address your remaining concerns below. Please inform us if you have any remaining concerns.

Limited theoretical analysis of when/why temporal alignment enables compositional generalization

We have added theoretical justification for why temporal alignment enables task compositionality to section 3.4, ''Temporal alignment and Compositionality.'' Intuitively, if goal representations are aligned with the previous states that lead to them, then goals which are slightly beyond the horizon of in-distribution tasks have a high probability of being in in-distribution (assuming the contrastive loss objective is minimized; see the section and the appendix for the full formalism). The key result is Theorem 1, which bounds the compositional out-of-distribution error of the policy in terms of the in-distribution error based on $\alpha=H'/H$, or how much larger the horizon $H'$ of the composed tasks is relative to $H$, the training horizon. We check this bound is meaningful by plotting it against a naive ''worst-case'' compositional generalization bound, showing it is meaningfully tighter in the new Figure 7. The result can be extended to compositionally-OOD language in addition to goals (Corollary 1).

How sensitive is the method to the choice of discount factor $\gamma$ in the temporal alignment objective? Was any ablation done on this hyperparameter?

In our preliminary experiments, we found there was not too much sensitivity to the precise value of $\gamma$, as long as $\frac{1}{1-\gamma}$ was proportional to the horizon $H$. This is supported by the new theoretical results which require $\frac{1}{1-\gamma}>H$ for the benefits in compositional generalization.

For the semantic generalization experiments (Scene C), how robust is the method to variations in object appearance beyond what was seen in training?

Our contribution is to study how structured representation learning objectives can enable compositional generalization. The related notion of object-level generalization is a distinct property, depending on the model architecture (ResNet), the diversity of the dataset (Bridge v2), and the generalization capabilities of the pre-trained visual embeddings (CLIP). We observe some robustness over object appearance, in line with the observations of past work in the Bridge v2 setting that has focused on other forms of generalization (e.g., SuSIE [1], GRIF [2], RT-X [3]).

How does the method scale to more complex real-world scenarios with greater environmental variation and uncertainty?

We have added discussion to our ''Limitations and future work'' section on how future work could explore similar compositionality properties in more complex environments such as bimanual, multi-agent, or cross-embodiment settings. Note that the tasks explored in this paper are of comparable difficulty to past works published at ICLR that use the BridgeDataset (v2) setting [1,4].

Dear Reviewer, Do you mind letting the authors know if their rebuttal has addressed your concerns and questions? Thanks! -AC

Need stronger justification for why TRA is preferable to VLM/LLM-based decomposition approaches

Could you provide quantitative comparisons with methods that use LLMs for task decomposition combined with foundation models (like Octo/OpenVLA) for execution?

What advantages does TRA offer over LLM-based decomposition approaches for long-horizon tasks? Please provide concrete examples and experimental results.

There have been many works in recent years that have used pipelines incorporating VLM/LLMs with low-level control policies to enable long-horizon behaviors [5,6,7,8,9]. The advantage of end-to-end approaches like TRA is they avoid compounding errors from multi-stage pipelines, and do not depend on large pre-trained models that need to be fine-tuned and prompted correctly. Note that representation-learning approaches like TRA still have access to some of general knowledge from VLM training by using pre-trained language and vision representations, but do not suffer the costs of hierarchical decomposition in doing so.

Since our focus in this paper is how the structure of successor representations can enable improved end-to-end compositional reasoning, we restricted our comparisons to be across end-to-end approaches. For quantitative results with VLM decomposition, we would like to refer to the evaluations in PALO [7], which feature some similar long-horizon tasks within the Bridge v2 dataset. While PALO achieves impressive long-horizon performance, it requires significant prompt engineering ([7], Appendix F) and expert demonstrations of all the new tasks. Similar limitations are seen with methods like RT-H [9], Inner Monologue [6], and SayCan [5].

In theory, VLM/LLM planners could be combined with TRA to solve even longer tasks when they are available. We have added discussion to the ''Limitations and related work'' section on these directions. Does this resolve your concern? Please let us know if there are any further comparisons we could provide.

References

[1] Black, K. et al., 2024. ''Zero-Shot Robotic Manipulation With Pre-Trained Image-Editing Diffusion Models.'' ICLR

[2] Myers, V. et al., 2023. ''Goal Representations for Instruction Following: A Semi-Supervised Language Interface to Control.'' CoRL

[3] O'Neill, A. et al., 2024. ''Open X-Embodiment: Robotic Learning Datasets and RT-X Models.'' ICRA

[4] Zheng, C. et al., 2024. ''Stabilizing Contrastive RL: Techniques for Robotic Goal Reaching From Offline Data.'' ICLR

[5] Ahn, M. et al., 2022. ''Do as I Can, Not as I Say: Grounding Language in Robotic Affordances.'' CoRL

[6] Huang, W. et al., 2022. ''Inner Monologue: Embodied Reasoning Through Planning With Language Models.'' CoRL

[7] Myers, V. et al., 2024. ''Policy Adaptation via Language Optimization: Decomposing Tasks for Few-Shot Imitation.'' CoRL

[8] Zawalski, M. et al., 2024. ''Robotic Control via Embodied Chain-of-Thought Reasoning.'' CoRL

[9] Belkhale, S. et al., 2024. ''RT-H: Action Hierarchies Using Language.'' arXiv:2403.01823

Thank you for your detailed and thoughtful review.

It seems your primary concerns relate to (1) a lack of theoretical justification for the method, (2) insufficient experimental evidence, and (3) issues with presentation. Please let us know if these changes address your concerns.

• For (1): we have added a new theory section (3.4), where we show how the TRA objective can induce a tighter bound on compositional generalization under some assumptions (see new Theorem 1)
• For (2): we have run additional ablations to show the robustness of TRA at real-world compositional generalization.
• For (3): we have made general improvements to writing with key changes highlighted in red based on your feedback.

No theoretical foundation for why temporal alignment should enable task compositionality is provided. Experiments on a single dataset do not provide enough empirical evidence to justify the claims of the paper without theoretical results.

We have added theoretical justification for why temporal alignment enables task compositionality to section 3.4, "Temporal alignment and Compositionality." Intuitively, if goal representations are aligned with the previous states that lead to them, then goals which are slightly beyond the horizon of in-distribution tasks have a high probability of being in in-distribution (assuming the contrastive loss objective is minimized; see the section and the appendix for the full formalism). The key result is Theorem 1, which bounds the compositional out-of-distribution error of the policy in terms of the in-distribution error based on $\alpha=H'/H$, or how much larger the horizon $H'$ of the composed tasks is relative to $H$, the training horizon. We check this bound is meaningful by plotting it against a naive "worst-case" compositional generalization bound, showing it is meaningfully tighter in the new Figure 7. The result can be extended to compositionally-OOD language in addition to goals (Corollary 1). Do these theoretical claims address your concern?

Loss function lacks clear explanation and intuition

We have added a section, "3.1 Motivation: Representations for Reaching Distant Goals," to provide intuition and explanation for the loss function. This intuition is then formalized in "3.4 Temporal alignment and Compositionality," where we show how using these representations can tighten the bound on the compositional generalization error under some assumptions (see new Theorem 1).

Additional experiments

We have run some additional ablations to show the robustness of TRA for real-world compositional generalization (see Table here)

How big is the computation overhead for the proposed method?

The computational overhead from TRA is relatively small ($<25\%$). The main contributor to this is the temporal alignment loss, since it requires forward and backward passes through the additional learned $\phi$ representation.

Performance of GRIF reported in the original paper for the same task is very different compared to the reported values here: "put the spoons on towels" from GRIF 0.9 and here "put the spoon on the towel" 0.2. How do you explain these big gaps?

Note the difference in the two tasks: "put the spoons on towels" (plural) v.s. "put the spoon on the towel." The plural version of the task requires compositional generalization as there are no examples of manipulating multiple spoons at the same time in the dataset. The large gap is precisely because GRIF does not enable better compositional generalization, unlike TRA.

Writing and clarity issues:

Thank you for noting these issues. We have substantially revised the text based on your comments and made general improvements to writing. Please let us know if any of the writing is still unclear.

Table 1 does not provide evidence for the made claims as just predicting trajectories without actual rollouts can be very misleading and robotics and does not have a clear correlation with success rate.

Yes, our main result is the substantial improvement in actual rollouts shown in Figure 2 and Tables 2 and 3. We included Table 1 to provide additional quantitative evidence for compositional generalization (in line with the definitions of generalization error in the new Theorem 1). If you would prefer, we could also cut Table 1 entirely or move it to the appendix. Would this address your concern?

Thank you for the response!

The request for adding a simulation based results on CALVIN or similar benchmark for reproducibility and future work still remain open.

Thank you for this suggestion. We believe that simulation results in a benchmark that explicitly tests compositional behaviors would be a valuable addition. In practice, many standard (offline) RL benchmarks actually fail to adequately test this property [1].

Since the time of our initial submission, a new offline RL benchmark, OGBench [2], has been released that explicitly tests compositionality and contains several manipulation tasks and datasets. We have implemented the TRA approach in this environment, and plan to include the full comparisons in a future revision.

Could you also elaborate in more detail the theoretical novelties and differences of TRA compared to similar prior work like VIP, LIV or Inference via Interpolation that also employ contrastive representation learning in the policy latent space? How does your loss formulation differ from these prior work?

There are several key differences between TRA and past representation learning approaches for robotics. At a high-level, TRA focuses on the question of compositionality. We show that a simple and scalable time-contrastive representation learning objective can produce this property when used for task representation, without requiring any explicit planning or reinforcement learning. In contrast, prior approaches like VIP and Inference via Interpolation focus on learning state representations to improve other aspects of decision making. We briefly discuss some of the technical distinctions relative to the mentioned approaches below.

VIP:

VIP [3] focuses on learning representations such that distances in the representation space correspond to goal-conditioned value functions. This work makes several assumptions about the environment: (1) that dynamics are fully deterministic, and (2) that goal-conditioned values are symmetric (i.e., $V(s;g)=V(g;s)$). Note that both of these assumptions are violated in real-settings: (1) real world dynamics are noisy/unpredictable, and (2) any external force acting on the system (e.g., gravity) breaks symmetry of goal-conditioned value functions—this is most dramatic with actions like opening the gripper, which at an elevation can drop an object, which is much easier to do than undo.

Under these assumptions, the VIP objective is a Fenchel dual of a time-contrastive representation objective. TRA does not require either of these assumptions, which may help with long-horizon tasks that are compositionally OOD and require awareness of the directionality of time and uncertainty.

LIV:

LIV [4] is an extension of the VIP objective which enables interoperability with language with an additional vision-language contrastive alignment loss. In practice this is implemented similarly to VIP, with the addition of a CLIP infoNCE [5] alignment loss. Similar approaches to connecting language and visual representations have been used in several recent works [6,7,8].

Inference via Interpolation:

Inference via Interpolation [9] shows how a form of time-contrastive representation learning can enable analytic, explicit subgoal inference by ''interpolating'' in the representation space. This requires several geometric assumptions: uniformity of the representation marginals [10] and the assumption that the state and goal encoders differ by a linear transform. The advantage of TRA is that we don't need to explicitly infer subgoals to compose behaviors, which avoids computational cost and compounding errors.

We will revise the paper to highlight these differences in our related work section.

Thank you for your review and suggestions.

It seems your primary concerns have to do with the evaluation tasks and baselines. We have run additional ablations on object-level reasoning as suggested. We discuss your remaining concerns below, highlighting revisions in the text with red. Do these changes fully address your concerns? We look forward to continuing the discussion.

No object-level ablation

We have run additional object-level ablations, including the particular task suggested by the author.

Table: object-level ablation tasks

The success rates for the new, reordered versions of the tasks do not differ from the original tasks significantly, showing the model is not just overfitting / memorizing sequences of events from the training data.

poor Baselines: Why not choose diffusion policies for imitation learning baselines?

The precise parameterization of our policy head is fairly orthogonal to the main contribution of our work, which is studying how the task representations enable compositionality. We picked our current ResNet (+FiLM conditioning [1]) backbone to be consistent with past work using the BridgeData v2 [2], which has found similar architectures to be most effective.

We have added discussion to our “Limitations and future work” section to address this.

too similar to VIP, why not give more explanations?

Indeed, there are tight connections between our work and past approaches that use self-supervised representation learning for robotics. We would like to emphasize that our representations and analysis tackle the problem of compositionality, which is distinct from the benefits to representation learning studied in papers like VIP [3], LIV [4], or Voltron [5]. We have added more discussion on these methods to the related work section.

Limited evaluation tasks … more complex or multi-agent settings might challenge TRA's capabilities

We agree that more complex settings are an exciting future direction. Our primary contribution in this paper is to show that a structured representation learning approach can provide an end-to-end robot learning algorithm for enabling a new property (compositionality) when used with an existing standard robot learning dataset (Bridge v2). We believe the scope of current evaluations on new robot manipulation tasks within the Bridge scenes is in line with prior work published at ICLR in similar settings [6,7]. We have added discussion to the future work section to suggest these future areas of exploration.

Dependence on Goal Representation Quality: Success in tasks depends heavily on the quality and specificity of goal representations, which may require fine-tuning for certain task types

We have added discussion on this limitation to the related work section, which is also a practical limitation for any representation learning objective. Does this address your concern?

It seems the font of the paper is weird. Have you chosen the right style?

We were using a computer modern font (MLModern) since the ICLR style files actually don’t specify a font. The revision has switched to Times New Roman. Does this address your concern?

References

[1] Perez, E. et al., 2018. "FiLM: Visual Reasoning With a General Conditioning Layer." AAAI

[2] Walke, H. et al., 2023. "BridgeData V2: A Dataset for Robot Learning at Scale." CoRL

[3] Ma, Y. et al., 2023. "VIP: Towards Universal Visual Reward and Representation via Value-Implicit Pre-Training." ICLR

[4] Ma, Y. et al., 2023. "LIV: Language-Image Representations and Rewards for Robotic Control." ICML

[5] Karamcheti, S. et al., 2023. "Language-Driven Representation Learning for Robotics." ICRA

[6] Black, K. et al., 2024. "Zero-Shot Robotic Manipulation With Pre-Trained Image-Editing Diffusion Models." ICLR

[7] Zheng, C. et al., 2024. "Stabilizing Contrastive RL: Techniques for Robotic Goal Reaching From Offline Data." ICLR

Dear Reviewer, Do you mind letting the authors know if their rebuttal has addressed any of your concerns and questions? Thanks!