Thank you for reviewing our paper. Before answering your questions, we would like to briefly highlight our contributions.

Our main contribution is to identify conditions under which low-rank structure naturally emerges in zero-shot reinforcement learning. Unlike prior work that assumes universal low-dimensional representations that generalize across tasks, e.g. Forward-Backward model (Touati and Ollivier, 2021; Touati et al., 2023) or goal-conditioned RL frameworks (Andrychowicz et al., 2017; Eysenbach et al., 2022b), we provide a theoretical basis for when and why such structure arises.

Specifically, we show that learning successor measures directly is statistically challenging, and demonstrate that a simple temporal shift of successor measures naturally induces low-dimensional representations.

Furthermore, we present preliminary numerical results in section 6 and additional experiments in Appendix H that illustrate the benefits of learning shifted successor measures. We believe that our work lays the foundation for practical algorithms for policy learning via shifted successor measures.

Motivation for using low-rank successor representations.

Thanks for this question: the objective of our paper is really to answer it! Without any structure, learning cannot be statistically efficient. That is why (and we explain this in the intro) many successful existing approaches in RL assume the existence of a low-rank structure, usually by assuming low-dimensional representation vectors $\phi(s,a)$ and $\psi(g)$. In this paper, we explain how such a structure appears naturally and universally (section 5), and we exemplify how the statistical efficiency of RL estimation procedure is improved by leveraging this structure (section 4). Note however that our results in Theorem 1 are always valid (not only for low-rank) and yet show the impact of having low-rank structures.

Connection between shifted-SR and optimal policy.

We thank the reviewer for raising this important question. There are three crucial points in addressing this issue:

(a) On the infeasibility of estimating $M_\pi$. If the reward function is known, the matrix product $M_{\pi} R$ directly gives the Q-function. Therefore, estimating $M_{\pi}$ allows policy evaluation for all rewards, which is of prime interest for transfer learning, multi-goal RL, etc. However, as discussed earlier, we argue that accurately estimating $M_{\pi}$ is generally infeasible due to the absence of a low-rank structure, and propose shifted successor measures as a way to circumvent this issue. For a moderate shift, $M_{\pi,k}$ remains reasonably close to $M_{\pi}$, while being more easily recoverable using low-rank approximation.

(b) Shifted-SR recover (nearly)-optimal policy in Goal Conditioned RL. In sparse reward settings, the most challenging in practice, policy updates based on shifted objectives yield policies that coincide with the optimal policies except when very close to the goal. Intuitively, for a fixed shift parameter $k$, if the goal is more than $k$ steps away from the current state, the Q-values computed using $M_{\pi,k}R$ are identical to those computed from $M_{\pi}R$. This ensures that the resulting policy is optimal until the agent enters the $k$-neighborhood of the goal, where exact optimality cannot be guaranteed. Nonetheless, as shown empirically in section 6 (and Appendix H), shifted-SR still recovers nearly optimal policies in practice.

(c) On the importance of entrywise guarantees for good policy recovery. Our work focuses on policy evaluation, and in this context, the entrywise guarantees we provide i.e. bounds $\Vert M - \widehat{M}\Vert_{2\to\infty}$ are crucial. Such guarantees imply that greedy policy $\pi_{\widehat{M}}$ derived from $\widehat{M}$ will perform well. This contrasts with weaker guarantees in global norms such as the Frobenius norm, which do not provide the same level of assurance for individual states. Thus, the strong entrywise bounds we establish are central to ensuring that accurate estimation of $M$ leads to high-quality policies.

Dear reviewer,

We hope our rebuttal has addressed your questions. We would greatly appreciate it if you could let us know whether there are any points you would like to discuss further or would like us to clarify. Thank you very much for your time and consideration.

Best regards,

The authors.

Thank you for reviewing our paper. Before answering your questions, we would like to briefly highlight our contributions.

Our main contribution is to identify conditions under which low-rank structure naturally emerges in zero-shot reinforcement learning. Unlike prior work that assumes universal low-dimensional representations that generalize across tasks, e.g. Forward-Backward model (Touati and Ollivier, 2021; Touati et al., 2023) or goal-conditioned RL frameworks (Andrychowicz et al., 2017; Eysenbach et al., 2022b), we provide a theoretical basis for when and why such structure arises.

Specifically, we show that learning successor measures directly is statistically challenging, and demonstrate that a simple temporal shift of successor measures naturally induces low-dimensional representations.

Furthermore, we present preliminary numerical results in section 6 and additional experiments in Appendix H that illustrate the benefits of learning shifted successor measures. We believe that our work lays the foundation for practical algorithms for policy learning via shifted successor measures.

On the scope and complexity of experimental validation.

As noted in the main text, Appendix H provides additional results on goal-conditioned tasks, which are discretized versions of standard (continuous) benchmarks for navigation tasks. There, we also discussed the impact of different data-collection policies and learning algorithms. Moreover, Appendix H.5 discusses extensions to non-tabular environments.

We agree that evaluating on high-dimensional environments is a valuable future direction. However, the goal of our experiments is to validate the theoretical insights under controlled conditions, where factors can be isolated and analyzed rigorously. We believe the current experiments, together with our theoretical contributions, make the paper comprehensive and self-contained.

Comparison with the baselines.

This paper introduces a novel insight not previously reported - assuming low-dimensional successor representations apriori is unnecessary, as low-dimensional structure naturally emerges when shifting successor measures. We believe this theoretical contribution is significant on its own and can lead to further advances.

Our algorithm is not directly comparable to standard baselines because it relies on SVD-based estimators, which are tailored to the tabular setting and do not trivially extend to high-dimensional environments. As discussed in Appendix H.5, extending the approach is possible but would likely come at the cost of the theoretical guarantees established in this work.

Finally, our experiments are designed to illustrate and validate the theoretical results in a controlled setting rather than to benchmark against existing methods, ensuring a clear link between theory and empirical results.

On the computational efficiency of shifted estimator.

In the tabular setting, the computational complexity of our algorithm is dominated by the singular value decomposition step. For more general settings discussed in Appendix H.5, the efficiency depends on the underlying algorithm for learning the successor measure. In some cases, such as when using contrastive learning approaches, incorporating the shift introduce minimal overhead and does not increase computational complexity.

Thank you for reviewing our paper. Before answering your questions, we would like to briefly highlight our contributions.

Our main contribution is to identify conditions under which low-rank structure naturally emerges in zero-shot reinforcement learning. Unlike prior work that assumes universal low-dimensional representations that generalize across tasks, e.g. Forward-Backward model (Touati and Ollivier, 2021; Touati et al., 2023) or goal-conditioned RL frameworks (Andrychowicz et al., 2017; Eysenbach et al., 2022b), we provide a theoretical basis for when and why such structure arises.

Specifically, we show that learning successor measures directly is statistically challenging, and demonstrate that a simple temporal shift of successor measures naturally induces low-dimensional representations.

Furthermore, we present preliminary numerical results in section 6 and additional experiments in Appendix H that illustrate the benefits of learning shifted successor measures. We believe that our work lays the foundation for practical algorithms for policy learning via shifted successor measures.

On the scope and complexity of experimental validation.

As noted in the main text, Appendix H provides additional results on goal-conditioned tasks, which are discretized versions of standard (continuous) benchmarks for navigation tasks. There, we also discussed the impact of different data-collection policies and learning algorithms. Moreover, Appendix H.5 discusses extensions to non-tabular environments.

We agree that evaluating on high-dimensional environments is a valuable future direction. However, the goal of our experiments is to validate the theoretical insights under controlled conditions, where factors can be isolated and analyzed rigorously. We believe the current experiments, together with our theoretical contributions, make the paper comprehensive and self-contained.

Restrictiveness of assumptions in Corollary 1.

The validity of the assumptions depend on the intrinsic dynamics but also on the choice of the policy. Overall, Corollary 1 requires the existence of an underlying structure (e.g. to have $\sigma_1, \ldots \sigma_r$ of order $1$, and $\sigma_{r+1}$ small), a sufficient shift (to have $\xi(M_{\pi,k})$ bounded), and a reasonably exploring policy (to have an underlying measure essentially invariant and uniform). We discuss these more in detail in Appendix C.1. The kind of setup where these apply well is for instance maze environments, like the four-room environment we give as an example in Theorem 4.

Extending the theory to continuous environments.

The notion of successor measure naturally extends to continuous state/action spaces using more general linear operators and measures instead of matrices (as is done in the forward-backward papers). We could extend Theorem 1 to this setting by discretization and leveraging smoothness properties of the dynamics (thereby requiring another assumption), in the spirit of the papers [1,2,3]. We restrict to the discrete setting first for simplicity, secondly because continuous settings are to be discretized anyway.

[1] Devavrat Shah, Dogyoon Song, Zhi Xu, and Yuzhe Yang. Sample efficient reinforcement learning via low-rank matrix estimation. NeurIPS 2020.

[2] Stefan Stojanovic, Yassir Jedra, and Alexandre Proutiere. Model-free low-rank reinforcement learning via leveraged entry-wise matrix estimation. NeurIPS 2024.

[3] Jiaming Xu. Rates of convergence of spectral methods for graphon estimation, ICML 2018

On the significance of relaxed accuracy.

In many scenarios, reaching the exact goal is unnecessary - being close to the goal is sufficient. Relaxed accuracy captures this notion. As shown in Figure 4 in section 6, shifted successor measures not only improve exact goal-reaching with low-rank approximation (c), but also perform better at reaching the neighborhood of the goal (d). This demonstrates the robustness of shifted successor measures - even when they do not yield the optimal policy, they still outperform the non-shifted baseline in getting close to the goal.

Policies with higher shifts perform relatively better because they are better aligned with relaxed objectives: if we consider a $k$-shift and thus neglect to collect the reward in the first $k$ steps, this matters at most within a radius $k$ around the reward, where a sub-optimal trajectory may not be penalized. By looking at relaxed accuracy, we do not account for this sub-optimality.

Intuition behind why shifting induces low-rank structure.

Shifting corresponds to multiplying by $P_{\pi}^k$, hence the shifted SM is the discounted sum of matrices $P_{\pi}^{k+l}$ for $l\ge 0$ which are in general more approximately low rank for large $k$: on the one hand, $I$ or $P_{\pi}$ present in the non-shifted SM have no reason to be low-rank, while $P_{\pi}^{k}$ converges to a rank $1$ matrix as $k \rightarrow \infty$ in the ergodic case (one eigenvalue is equal to $1$ and all others are strictly below $1$ in absolute value). The shift is a way to interpolate between these two extreme cases. This is the mixing intuition we make more rigorous in section 5.

Furthermore, this intuition works even better in more stochastic problems, since stochastic dynamics are more likely to be ergodic. Ergodicity may not be required however, provided we restrict to an appropriate sub-dynamics: see our discussion in Appendix C.1.

On the limitations of theoretical results.

Thanks for this remark. Regarding limitations, we discuss briefly the different terms of our upper bound after Theorem 1, while a more complete discussion is also given in Appendix C. We will highlight that more and summarize in the camera-ready version of the paper.

On the limitations of experimental results.

We thank the reviewer for this comment and refer them to Appendix H for a detailed description of the experimental setup. In the camera-ready version, we will add a paragraph in section 6 explicitly linking the experimental assumptions and design to the theoretical analysis.

Dear reviewer,

We hope our rebuttal has addressed your questions. We would greatly appreciate it if you could let us know whether there are any points you would like to discuss further or would like us to clarify. Thank you very much for your time and consideration.

Best regards,

The authors.

Thank you for reviewing our paper. Before answering your questions, we would like to briefly highlight our contributions.

Our main contribution is to identify conditions under which low-rank structure naturally emerges in zero-shot reinforcement learning. Unlike prior work that assumes universal low-dimensional representations that generalize across tasks, e.g. Forward-Backward model (Touati and Ollivier, 2021; Touati et al., 2023) or goal-conditioned RL frameworks (Andrychowicz et al., 2017; Eysenbach et al., 2022b), we provide a theoretical basis for when and why such structure arises.

Specifically, we show that learning successor measures directly is statistically challenging, and demonstrate that a simple temporal shift of successor measures naturally induces low-dimensional representations.

Furthermore, we present preliminary numerical results in section 6 and additional experiments in Appendix H that illustrate the benefits of learning shifted successor measures. We believe that our work lays the foundation for practical algorithms for policy learning via shifted successor measures.

Choosing the shift parameter $k$ in practice.

As discussed in section 5, the optimal shift value $k$ can be derived in certain settings. For general environments with complex dynamics, $k$ must be treated as a tunable hyperparameter. Its value is related to other hyperparameters, such as the discount factor. This is consistent with prior work - for instance, HIQL (Park et al., 2024) treats the number of subgoal steps $k$ as a hyperparameter, tuning it per environment. We argue that selecting shift parameter $k$ in our setting is no more difficult than setting the subgoal step parameter in HIQL.

Practicality of the assumptions in section 5.

As far as we are aware, we are first to analyze what happens when shifting successor measures, and we believe this idea is novel and interesting for further practical analysis. Our main result of section 4 makes no assumption and applies to arbitrarily complex, discrete dynamics. It could be extended to continuous settings, with the right regularity assumptions, as we answered Reviewer tYCN. The ideas of section 5 are also universal: mixing phenomena exist beyond discrete Markov chains, which is mainly a way to avoid more tedious formalism in continuous setting. Furthermore, these ideas could apply even to Markov processes which do not mix in the classical sense. We discuss some extensions in Appendix C.1. Thus we believe our analysis holds in good generality and provides insight that might be interesting for their own sake. What we do in section 5.2 (decomposable Markov chains) is to just provide an example where the analysis can be directly applied. There, we indeed make some assumptions.

On the scope and complexity of experimental validation.

As noted in the main text, Appendix H provides additional results on goal-conditioned tasks, which are discretized versions of standard (continuous) benchmarks for navigation tasks. There, we also discussed the impact of different data-collection policies and learning algorithms. Moreover, Appendix H.5 discusses extensions to non-tabular environments.

We agree that evaluating on high-dimensional environments is a valuable future direction. However, the goal of our experiments is to validate the theoretical insights under controlled conditions, where factors can be isolated and analyzed rigorously. We believe the current experiments, together with our theoretical contributions, make the paper comprehensive and self-contained.