Horizon Generalization in Reinforcement Learning

Dear reviewer,

Thank you for your responses and suggestions for improving the paper.

We have revised the text (changes highlighted in red) to improve the clarity of our theoretical claims, and clearly state the practical implications of our theory and experiments for GCRL algorithms throughout the paper. Regarding your concern about real-world applications and limitations of these generalizations in more complex settings, we will evaluate a higher dimensional setting (humanoid) with a more challenging control component. We hope that these experiments will sufficiently address any remaining concerns.

I am also concerned about real-world applications (e.g., more complicated tasks). In real-world settings, the planning invariance and horizon generalization seem to be limited in their application to general cases, although they indeed exist as mentioned by the authors. In the high-dimensional settings (Section 6.2), it is hard to derive insightful results.

As mentioned in Section 4.5, (1) the optimality over local trajectories and the (2) coverage of local trajectories are both concerns. For (1), there are unavoidable sources of errors from function approximation and inherent noise in training. To push on this potential weakness, we will further test the presence and relationship between planning invariance and horizon generalization settings in a higher d.o.f. settings like Humanoid [1].

We would also like to highlight that we have already tested planning invariance, horizon generalization, and their relationship in Ant, which is a standard, 27-dim observation 8DoF control, RL benchmark used as a primary evaluation in many recent prior RL work [3,4,5]. Since the initial submission, we have extended the results in Figure 5 to show further horizon generalization in this setting.

References

[1] Tunyasuvunakool, S. et al., 2020. "dm_control: Software and Tasks for Continuous Control." Software Impacts

[2] Park, S. et al., 2024. "OGBench: Benchmarking Offline Goal-Conditioned RL." arXiv:2410.20092

[3] Zheng, C. et al., 2023. "Contrastive Difference Predictive Coding." ICLR

[4] Fujimoto, S. et al., 2018. "Addressing Function Approximation Error in Actor-Critic Methods." ICML

[5] Freeman, C. et al., 2021. "Brax—a Differentiable Physics Engine for Large Scale Rigid Body Simulation." NeurIPS Datasets and Benchmarks

Dear Reviewer,

Thank you for your responses and suggestions.

It seems your main suggestion is to provide more intuition for the theoretical claims, including a discussion of the relationship with prior work on representations in RL. We have made revisions in the paper to address these concerns and believe that the paper motivation is now much clearer. Together with the discussion below, do these fully address your concerns?

We have revised to the paper to provide more intuition and background:

1. Highlight relationships with prior state abstraction works such as bisimulation: In Section 2, we now compare and contrast horizon generalization with other forms of generalization captured by prior state abstraction methods. We highlight that, to our knowledge, there has not been prior work directly addressing generalization from short to long horizons. Furthermore, planning invariant policies map state-action-waypoint pairs and state-action-goal pairs to the same latents, despite different associated rewards and Q-values (see Section 3 for relation between quasimetric and underlying reward function in an MDP). This is already a violation of the Q-irrelevance relation and the stronger bisimulation relation [1]: the invariances captured by planning invariance and such state abstractions are fundamentally different.
2. Building intuition for planning invariance: We have updated our planning invariance figure to help build intuition and highlight Section 5: Methods for Planning Invariance: Old and New.
3. Takeaways: We have highlighted the high-level takeaways for the reader throughout the paper (see end of Section 1, Section 2, end of Section 4.4, and Section 7). By theoretically and empirically linking planning with horizon generalization, our work suggests practical ways (i.e., quasimetric methods) to achieve powerful notions of generalization from short to long horizons. Do these revisions, together with the discussion below, fully address your concerns? We look forward to continuing the discussion!

There doesn't seem to be any suggestions on how to induce this kind of invariance.

This invariance is precisely the structure imposed by quasimetric architectures—we have clarified the notation so this is now clearly stated as Theorem 2 in the text. We have also added a short summary of the concrete takeaways to the conclusion section, and a summary of different architectures and losses and how they relate to the planning invariance of quasimetric architectures. Additional discussion is presented Section 5: “Methods for Planning Invariance, Old and New” as well as Appendix C. For example, for the forward perspective, we can induce planning invariance (1) explicitly via dynamic programming or (2) implicitly using a quasimetric architecture. We empirically demonstrate that quasimetric architectures, when combined with a contrastive metric learning approach [3], are relatively planning invariant compared to other architectures (see Figure 5, right).

relation to compositional generalization in language-conditioned RL

Indeed, a key advantage of language as a form of task specification is its ability to easily express compositional relationships. Future work could investigate how planning-invariant reasoning could be extended to these alternate task modalities, through, e.g., learned mappings from the task space into the quasimetric goal space [2]. We have added discussion to our conclusion (marked in red, under “Limitations and future work”) on how future work could tackle these challenges.

give more examples of non-trivial cases where planning invariance is or is not met, in order to build intuitions

We have replaced the figure in Section 4.1 with Figure 2. With this revised figure, we can (1) visualize different 2D translation actions for different navigation goals (purple star and brown star) as well as optimal (tan) and suboptimal (gray) trajectories, (2) identify that the left policy has no planning invariance (policies towards goal and waypoint are not the same, where the policy selects a suboptimal action while navigating towards the goal and optimal action while navigating towards the waypoint) while the right policy has planning invariance (indicated by the same, optimal policy whether directed towards a goal or the waypoint), and (3) concretely connect the presented maze task to actual sequential decision-making problems like navigation with obstacles.

Dear Reviewer,

Thank you for your responses and suggestions. It seems like your main concerns have to do with notation and clarity, as well as how the theory could relate to practical new methods. We have made revisions in the paper to address these concerns and believe that the paper clarity is now much higher. We hope that these improvements have increased the readability of the main theorems and make the takeaways clear. Together with the discussion below, does this fully address your concerns? We look forward to continuing the discussion!

Clarifying notation.

We have made several revisions to the paper to improve clarity:

1. We clearly define the successor distance over actions $d(s,a,g)$, which, using the setup in [1], is $\min\_{\pi} \frac{p^{\pi}\_{\gamma}(s\_K = g|s\_{0} = g)}{p^{\pi}\_{\gamma}(s\_K = g|s\_{0} = s,a)}$ (Equation 5) where $p^{\pi}\_{\gamma}(s\_K = g|s\_{0} = s)$ is the discounted state-occupancy measure (Equation 4) and $p^{\pi}\_{\gamma}(s\_K = g|s\_{0} = s, a)$ is the discounted state-occupancy measure with actions (Equation 5). This is the temporal distance to the goal $g$ conditioned on starting at state $s$ and taking action $a$ under an optimal policy.
2. We note that the minimization of $d(s,a,g)$ over actions $a$ is used for action selection: Myers et al. [1] show that $\text{arg}\min_a d(s,a,g)$ corresponds to $\text{arg} \max\_a Q(s,a,g)$ of the Q-function with respect to the goal-conditioned RL MDP.
3. We highlight in Sections 3 and 4.2 that the reward for the MDP using the successor distance is $r(s) = \delta\_{(s,g)}$, which is a Kronecker delta function that evaluates to 1 at the goal and 0 at intermediate states.
4. We have revised notation in Section 4 and Appendix C to use more standard math terms. For example, we now use set notation to construct $\text{arg} \min$ rather than introduce a “symmetry-breaking” $\text{arg} \min$ over equally optimal actions and waypoints. We have also renamed certain objects for clarity (i.e. policy $\pi^f\Rightarrow \pi^{\text{Fix}}$ and planning operator $\text{Plan}^{f} \Rightarrow \text{Plan}^{\text{Fix}}$ for the deterministic, “fixed” goal, controlled setting).

The changes to the proofs for clarity have been highlighted in red. Please let us know if additional clarifications will be helpful.

How does the paper inspire future algorithm development?

We have added a new discussion in Appendix D about how our theoretical results might inspire future algorithm development, as well as a summary of the concrete takeaways in the conclusion, and a comparison of existing GCRL critic parameterizations and how they relate to quasimetric architectures (Table 1). We agree that the primary takeaway is not a new method, but rather (1) concrete definitions of planning invariance and horizon generalization, (2) proofs linking quasimetrics, planning invariance, and horizon generalization, and (3) experiments showing that horizon generalization is feasible and nontrivial. That said, our experiments show that the modified CMD algorithm [1] combined with the backward infoNCE [4] loss provides a good base algorithm that empirically shows some of these properties while being motivated by our theory. We believe these results are significant and in line with prior ICLR works that highlight a phenomenon/property of existing methods [2, 3].

Typos, use of “lemma”

Thank you for these suggestions. We have fixed the typos and added definitions for the terms you noted, and additionally renamed the core (theoretical) results as “theorems” instead of “lemmas.”

References

[1] Myers, V. et al., 2024. "Learning Temporal Distances: Contrastive Successor Features Can Provide a Metric Structure for Decision-Making." ICML

[2] Richens, J. et al., 2024. "Robust Agents Learn Causal World Models." ICLR

[3] Subramani, R. et al., 2024. "On the Expressivity of Objective-Specification Formalisms in Reinforcement Learning." ICLR

[4] Bortkiewicz, M. et al., 2024. "Accelerating Goal-Conditioned RL Algorithms and Research." arXiv:2408.11052

Thank you for your response.

Thanks for the update for the undefined symbols. This means that my original assessment is correct, and I didn't miss anything. I also noticed that the definitions of several key symbols, which are used in major results such as Theorem 1, 2, and Definition 3, are missing in the original submission. I'm curious as to why they were omitted.

Due to space constraints, we initially moved some definitions to the appendix and omitted some standard definitions from prior work (for instance, the discounted state occupancy $p\_\gamma^\pi$ notation you noticed is used in numerous prior works [1,2,3,4,5]; similarly the goal-conditioned reward notation is defined in [6,7,8]).

We now include these definitions in the main text to make the paper more self-contained. We hope the current revision addresses this concern.

Additionally, the authors did not address my concerns regarding the confusing use of the terms ''metric'' and ''quasimetric.'' This could make the paper difficult for readers to understand.

The formal construction of the distances in this paper all rely on the quasimetric definition, which relaxes the more familiar notion of a metric by allowing asymmetric distances. While this notion of asymmetry is important in many MDPs, the key benefit of viewing things as (quasi)metrics is the triangle inequality, which is what enables planning invariance and horizon generalization. We use the term ''metric'' in some places when we wish to highlight the triangle inequality property rather than asymmetry, in line with prior work that formally relies on quasimetrics [9,10]. Note that since the definition of a metric is stronger than that of a quasimetric, all our results hold for metrics as well. We hope this note as well as our latest revision clarify this point.

Thank you for addressing how this work could inspire future research. It would be beneficial if the authors could present stronger results in the next version, as it would significantly enhance the paper's impact.

Thank you for this suggestion. We can include results in additional AntMaze layouts [11] and the Humanoid environment [12] in our next revision. Future work may also benefit from evaluation on the newly-released OGBench baseline [13], which features a diverse set of long-horizon tasks in the offline setting.

---

References

[1] Eysenbach, B. et al., 2021. ''C-Learning: Learning to Achieve Goals via Recursive Classification.'' ICLR

[2] Choi, J. et al., 2021. ''Variational Empowerment as Representation Learning for Goal-Conditioned Reinforcement Learning.'' ICML

[3] Myers, V. et al., 2024. ''Learning to Assist Humans Without Inferring Rewards.'' NeurIPS

[4] Schroecker, Y. and Isbell, C., 2020. ''Universal Value Density Estimation for Imitation Learning and Goal-Conditioned Reinforcement Learning.'' arXiv:2002.06473

[5] Hoang, C. et al., 2021. ''Successor Feature Landmarks for Long-Horizon Goal-Conditioned Reinforcement Learning.'' NeurIPS

[6] Fang, K. et al., 2022. ''Generalization With Lossy Affordances: Leveraging Broad Offline Data for Learning Visuomotor Tasks.'' CoRL

[7] Yang, R. et al., 2023. ''What Is Essential for Unseen Goal Generalization of Offline Goal-Conditioned RL.'' ICML

[8] Ghosh, D. et al., 2019. ''Learning Actionable Representations With Goal Conditioned Policies.'' ICLR

[9] Liu, B. et al., 2023. ''Metric Residual Network for Sample Efficient Goal-Conditioned Reinforcement Learning.'' AAAI

[10] Myers, V. et al., 2024. ''Learning Temporal Distances: Contrastive Successor Features Can Provide a Metric Structure for Decision-Making.'' ICML

[11] Fu, J. et al., 2021. ''D4RL: Datasets for Deep Data-Driven Reinforcement Learning.'' arXiv:2004.07219

[12] Gu, X. and Wang, Y., 2024. ''Advancing Humanoid Locomotion: Mastering Challenging Terrains With Denoising World Model Learning.'' RSS

[13] Park, S. et al., 2024. "OGBench: Benchmarking Offline Goal-Conditioned RL." arXiv:2410.20092

Dear Reviewer,

Thank you for your detailed review and suggestions. 

Your primary concerns seem to be (1) issues with the motivation and formalism for planning invariance and horizon generalization and (2) concerns about the utility of our formalism in practical environments.

We have made revisions based on your suggestions highlighted in red to clarify the theoretical results, added discussion on concrete takeaways to our conclusions, and updated our experiments to support our theoretical claims.

Do these changes fully address your concerns? We believe the current claims in the paper are correct (regarding “this is a bit far-fetched”) and have clarified the proofs in Appendix B. If you still believe any of the proofs or other writing is unclear, please inform us and we will be happy to make further revisions.

Generalization in RL and in goal-conditioned setting can come in different forms. While generalizing over different horizons is important, why do you think invariance will help here? 

In ML, there is a strong connection between generalization and invariance where enforcing invariance facilitates generalization. For example, neural networks invariant to rotation show better abilities to generalize over rotated test sets [1, 2, 3]. This body of work motivates our analysis; we have revised Section 4.1 to refer to such work. We expect generalization over different horizons in RL to also relate to a form of invariance, and identify planning invariance as a possible mechanism to achieve horizon generalization.

How does the method relate to when we are no longer looking at smaller horizons as being sub-goals of the longer horizons? 

In general, any goal reached at a given horizon will necessarily have possible shorter-horizon subgoal(s) just by taking a distribution of states reached ``on the way’’ to the goal. This is mathematically formalized through our induction argument, which necessarily assumes that you can use smaller problems to solve larger problems. Does this clarify your concern? We are not entirely sure if we have understood it correctly, but would be happy to provide further information if needed.

Or even when we are not considering the goal as a condition. Do you think this work can include those cases?

Outside the goal-conditioned setting, the distance metric in its current form loses meaning and we lose the benefits of ``reward-free” goal-conditioned reinforcement learning; we have added discussion of prior work and motivation for goal-conditioned RL in Section 2.  However, our result suggests that if you force your reward structure to take on values that lead to quasimetric Q-functions, or Q-functions that are scaled and shifted versions of quasimetrics, then we should expect horizon generalization in goal-free RL. How to enforce such constraints on reward functions in non goal-conditioned settings is still an open question, but could have potential implications on how to engineer rewards (that correspond to a desired objective, even if it isn’t a fixed goal) to achieve horizon generalization. 

The writing of the paper needs some work to strengthen the motivation and claims. The related works section can also benefit from further elaboration and coverage of the related literature.

As suggested, we have made several revisions to make the paper more clear and add motivation. We have (1) expanded our prior works to compare horizon generalization to other forms of generalization with latent space abstractions (see Section 2) and motivated the goal-conditioned setting, (2) clarified the limitations of planning invariance as a mechanism to achieve horizon generalization (Section 4.2 and 4.5), (3) added a remark to highlight that horizon generalization, as defined in this paper, is nontrivial (Remark 2), and (4) replaced the dominoes setting example with a more complex navigation task to build intuition for planning invariance (Figure 2). 

Experiments are over simpler environments / task settings.

We would like to note that the primary Ant environment studied involves 8DoF control and 29d observations [4], and is challenging enough to have been used as a primary evaluation for several recent RL works [5,6,7,8]. Since the submission, we have expanded the results in the Ant environment to study generalization to quantitatively more distant goals, and have additionally run preliminary experiments in the Humanoid [4] environment showing similar trends, which will be included in the final version of the paper along with the existing results.