Temporal Difference Flows

We thank the reviewer for the detailed and thoughtful evaluation, and for recognizing the clarity, rigor, and contributions of the work—both theoretical and empirical. Your careful reading and constructive suggestions are greatly valued.

Number of seeds and statistical reporting

Thank you for raising this important point. While we report results using 3 random seeds, this is done per policy and task, across multiple tasks within each domain—resulting in a substantial number of evaluations across the full benchmark suite. For generative modeling alone, the single- and multi-policy experiments involve over 800 individual runs, not including the computational cost of training the TD3 and FB policies. Additionally, in practice, we observed consistently low variance across runs, especially in value estimation and planning metrics, which informed our choice. We’ll be happy to include confidence intervals for figures in the revised manuscript.

Related work and comparison to GPI techniques

Thank you for the keen suggestion. While we did not evaluate the geometric switching variant of GPI proposed in [1], we see it as a natural and promising direction for future work—particularly given that TD-Flow significantly improves the underlying GHM used in their approach. We believe combining TD-Flow with geometric switching can further improve performance, but given the technical depth already required to develop and analyze TD-based flow and diffusion models, we felt this extension was best left to future work that can give it proper attention. A more focused related work discussion is already included in Appendix C, and we’d be happy to bring specific references or points into the main text if there are any the reviewer feels would benefit from greater emphasis.

Clarifications and presentation improvements

• Line 29: Thank you for catching this error, as defined in the appendix this should read geometric not generative.
• Line 59: We restrict to deterministic policies primarily for clarity and simplicity of exposition, especially in the theoretical sections. We’ll clarify this point in the revised manuscript.
• Equations 2 and 3: Thank you for pointing this out. We will add a definition of $\rho$ in the text. Specifically, $\rho$ denotes the empirical distribution over transitions sampled from a replay buffer or dataset.
• Many variants and method comparison: Thank you for the suggestion. To improve clarity, we’ll expand Table 3 in the Appendix to include the TD-diffusion variants and highlight the key differences between all methods. If space allows, we’ll consider moving this table to the main text; at the very least, we’ll ensure it’s clearly referenced from the main body to help guide readers.

---

We’re grateful for the reviewer’s thoughtful feedback and are excited to incorporate these refinements into the final version. We hope the improved clarity and added discussion further strengthen the case for the paper’s significance.

References

[1] Shantanu Thakoor, Mark Rowland, Diana Borsa, Will Dabney, Rémi Munos, and André Barreto. Generalised Policy Improvement with Geometric Policy Composition. International Conference on Machine Learning (ICML), 2022.

We sincerely thank the reviewer for their detailed and thoughtful evaluation. We appreciate the recognition of our theoretical contributions, the robustness of TD²-CFM and TD²-DD in long-horizon modeling, and the strong planning performance demonstrated through Generalized Policy Improvement (GPI).

Impact of dataset quality and scale on planning and modeling performance

We agree that understanding how data quality and scale affects performance is an important question. While we did not vary dataset size in this work, we note that all prior GHM methods similarly rely on offline data and suffer from instability and degraded performance even in much simpler and lower-dimensional environments than those we consider. Our key contribution is to show that by integrating temporal difference structure into both the learning objective and sampling process, we can significantly improve both stability and modeling performance. This is supported by both our theoretical analysis and strong empirical results demonstrating robustness across complex control tasks.

For the planning experiments, we use the same standard offline dataset as the Forward-Backward (FB) representation from [2], which has become a common benchmark in this setting. This choice enables a direct comparison and reveals a compelling result: TD²-based GHMs are able to discover significantly better policies from the same data. This highlights both the strength of our method and the fact that FB fails to identify the best policies within its own class—despite having access to the same data and that its reward embedding is purported to be optimal. We will clarify this in the revised manuscript to better contextualize our design choices.

Task diversity and inclusion of discrete/stochastic domains

Thank you for this thoughtful suggestion. In this paper, we focused on continuous control tasks as they are widely used benchmarks for evaluating generative models and the successor measure. Importantly, these settings already pose significant challenges for long-horizon modeling, and prior methods struggle to model even lower-dimensional versions of these tasks.

As a point of reference, the Ms. Pacman maze experiment in [1] only models the successor measure of 2D $(x,y)$ coordinates of the agent (the backward embedding takes as input $(x,y)$ and not full RGB frames), whereas our method learns to model much higher-dimensional state spaces—including full-body continuous dynamics. Additionally, our Pointmass maze tasks are sparse-reward and require long-horizon reasoning, making them meaningfully different in structure and challenge from the locomotion tasks.

We also agree that moving toward partially observable settings would be a valuable next step. At the same time, such environments introduce a distinct set of challenges, including belief-state estimation, memory, and uncertainty modeling, which we view as part of a larger and important research direction. We see our work as establishing a strong foundation for scalable, stable successor state modeling in fully observable settings, and we’re excited about the community building on this in future work.

Use of TD and potential training instability due to bootstrapping

Thank you for raising this important point. You're absolutely right that bootstrapping—especially when combined with function approximation and off-policy data—can introduce instability in reinforcement learning. These challenges are not fully avoidable in our setting either, as TD²-based GHMs retain all the characteristics of the so-called "deadly triad."

That said, a central goal of our approach is to mitigate these issues as effectively as possible. By integrating temporal-difference structure into flow and diffusion-based models, we observe significantly improved training stability and sample efficiency compared to prior GHM methods, which often struggle even at relatively short horizons. Our theoretical results further support this by showing that tighter coupling between sampling and regression reduces gradient variance, contributing to more stable learning.

In practice, we also incorporate common stabilization techniques—such as employing an exponential moving average target network. While instability remains a core challenge in TD-based methods, we believe our work represents a meaningful step toward more stable and scalable generative modeling of the successor measure.

Formatting

• L189: Thank you—we will correct the formatting issue in the final version.

---

We hope our responses help to contextualize the design choices and emphasize the broader relevance of our contributions.

References

[1] Ahmed Touati and Yann Ollivier. Learning One Representation to Optimize All Rewards. Neural Information Processing Systems (NeurIPS), 2021.

[2] Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does Zero-Shot Reinforcement Learning Exist? International Conference on Learning Representations (ICLR), 2023.

We thank the reviewer for the thoughtful and encouraging feedback. We're especially glad to hear that the paper was clear and accessible, even for readers less familiar with generative modeling.

What benefits does TD bring to flow/score matching?

This is a great question, and we appreciate you highlighting it. We’ll be sure to clarify it in the revised manuscript.

Temporal Difference (TD) learning, when applied to flow models, brings the classic benefits of TD methods into learning generative models of the successor measure of a given policy. By using current predictions to estimate future outcomes, TD allows us to train on shorter segments of data while still capturing long-horizon structure. More concretely, this provides several key advantages:

• Reduced variance and greater sample efficiency: Rather than relying on full long-horizon rollouts (which can be noisy, high variance, or unavailable), TD updates allow us to learn from short segments (e.g., 1-step transitions) and still reason about long-term behavior.
• “Stitching” across time: TD enables combining information from different parts of the environment/dataset to construct consistent long-horizon predictions, even when full trajectories are missing or limited.
• Off-policy learning: TD naturally supports learning from off-policy data—that is, trajectories collected from different policies or behaviours. This is critical in the offline setting, where the model must learn from a fixed dataset. Our planning results in particular are only possible because of this capability: TD-Flow can plan effectively using policies not seen during data collection.

What do Theorems 2 and 3 imply intuitively?

The theorems compare the variance in estimating the gradient of the loss functions. The smaller the variance, the less the number of data samples and generative samples are needed to have an accurate estimate of the gradient. Intuitively, we can think of the gap in variance as being caused by the misalignment of the conditional vector field $u_t$ with the true vector field $v_t$. For conditional paths, under certain conditions (i.e., coupling between $(X_0, X_1)$ and straight non-crossing paths) there is alignment and this gap is zero.

Why does Eq(1) hold?

By definition we have,

$ Q(s, a) &= \mathbb{E}_{\pi} \left[ \sum\_{k=1}^\infty \gamma^{k-1} r(S\_{t+k}) \mid S\_t = s, A_t = a \right] \\\\ &= \sum\_{k=1}^\infty \gamma^{k-1} \int\_{\mathcal{S}} r(x)\\, \mathrm{Pr}(S\_{t+k} = x \mid S\_t = s, A\_t = a; \pi) \mathrm{d}x \\\\ &= \int\_{\mathcal{S}} r(x) \underbrace{\sum\_{k=1}^\infty\\, \gamma^{k-1} \mathrm{Pr}(S\_{t+k} = x \mid S\_t = s, A\_t = a; \pi) \mathrm{d}x}\_{(1-\gamma)^{-1}m^\pi(x\mid s,a)} \\\\ &= (1-\gamma)^{-1} \mathbb{E}\_{X \sim m^\pi(\cdot\mid s, a)} \left[ r(X) \right] $

The MDP definition is different from the RL community, where γ acts on reward and value functions.

Thank you for pointing this out. You're absolutely right that in standard RL formulations, the discount factor $\gamma$ typically appears in the definition of return, i.e., as a weight on future rewards.

In our setting, since we model the successor measure directly, $\gamma$ instead appears in the distribution over time steps—specifically, it defines a geometric distribution over how far into the future we sample. This reparameterization is equivalent but shifts the emphasis from discounting rewards to discounting the probability of future state occurrences. Importantly, this does not imply that the MDP itself terminates with probability $1 – \gamma$ at each step—rather, it’s a modeling choice that reflects the probabilistic weighting of future events, not actual environment termination.

Since this interpretation doesn’t appear frequently in standard RL formulations, we’ll expand on it in the revised version to help clarify the connection.

---

We hope this response strengthens your overall impression of our work.

We thank the reviewer for their thoughtful and constructive feedback, and we’re glad that the novelty and potential of TD-Flow and its variants were appreciated. We respond below to the key points raised and will revise the paper accordingly to better highlight these aspects.

Comparison to direct Q-function learning and finite-horizon planning methods

We agree that placing TD-Flow in the context of existing value-based and model-based planning methods is important. Our approach offers several distinct advantages:

• Generalization across reward functions: Unlike traditional value-based methods, which are tied to a specific task, TD-Flow learns the successor measure which is reward agnostic. This enables zero-shot transfer to new reward functions, which we demonstrate in our planning experiments—an especially valuable property in offline or multi-task environments.
• Long-horizon reasoning: Because TD-Flow directly learns to sample from long-horizon state distributions, it avoids limitations of standard MPC methods (e.g., shooting or cross-entropy methods). This is particularly beneficial in sparse reward or goal-conditioned tasks, where short-horizon planners fail to make meaningful progress due to limited planning depth.
• Planning-time efficiency: Although training iterative generative models may be more expensive, inference can be significantly more efficient. For example, with a discount factor of 0.99, the expected time horizon is ~100 steps. Whereas a model-based planner would need to simulate 100 one-step transitions to reach that horizon, TD-Flow can sample directly from that distribution in far fewer steps using an ODE solver—often with just 10 integration steps—yielding substantial computational savings at inference time.

To directly assess this comparison, we include results for a Model Predictive Path Integral (MPPI) controller with a learned dynamics model. We train a similar capacity dynamics model to that of TD-Flow before evaluating MPPI with a finite horizon of 32 for locomotion tasks and 128 for maze, where at each step we sample 256 action candidates and perform 10 optimization rounds with 64 elites (top-k actions) per round. The results below show that TD-Flow significantly outperforms MPPI in 3/4 domains with comparable results in Walker. MPPI notably displayed instability related to compounding errors in environments with difficult to model dynamics. We hope this helps clarify the benefits of directly modeling the successor measure.

Run-time analysis and computational cost

This is a valuable point. Both the runtime at training and inference is dominated by the number of iterative steps to solve the ODE or SDE during sampling. A key insight—shown theoretically in Appendix E.7—is that our method naturally minimizes the transport cost which enables the use of a relatively small number of solver steps without significantly compromising accuracy.

To support this claim, we ran an empirical analysis showing how prediction quality degrades as we reduce the number of integration steps on the Loop task in Pointmass Maze. The results below show that TD²-CFM remains robust even at coarse discretizations of the ODE with as little as 5 integration steps, while we observewith a predictable degradation whenas the number of steps is too smalldecreases. Additionally, orthogonal work on consistency models [1] and self-distillation [2] can be applied to TD-Flow to further increase accuracy for few-step generation.

---

We thank the reviewer again for their insightful suggestions and willingness to reconsider their evaluation. We hope our response clarifies the unique strengths of TD-Flow and its potential to the community.

References

[1] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency Models. International Conference on Machine Learning (ICML), 2023.

[2] Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One Step Diffusion via Shortcut Models. International Conference on Learning Representations (ICLR), 2025.

[3] Manan Tomar, Philippe Hansen-Estruch, Philip Bachman, Alex Lamb, John Langford, Matthew E Taylor, and Sergey Levine. Video occupancy models. CoRR, abs/2407.09533, 2024.