Temporal Difference Flows

Temporal Difference Flows

Speakers: Jesse Farebrother, Matteo Pirotta, Andrea Tirinzoni, REMI MUNOS, Alessandro Lazaric, Ahmed Touati

Conference: ICML 2025

YouTube: https://slideslive.com/39043985

Overview

This article delves into "Temporal Difference Flows," a groundbreaking work presented at ICML 2025 by Jesse Farebrother, who conducted this research during an internship at Meta, alongside a team of distinguished collaborators including Matteo Pirotta, Andrea Tirinzoni, REMI MUNOS, Alessandro Lazaric, and Ahmed Touati. The talk introduces a novel approach to learning geometrically discounted horizon models, often referred to as gamma models or related to the successor representation in reinforcement learning. These models aim to predict future states directly, discounted by their temporal distance, offering a compelling alternative to traditional world models that often struggle with long-term predictions due to compounding errors.

The core contribution of this research lies in demonstrating that by leveraging the inherent structure of iterative generative models, specifically flow matching, it is possible to robustly and reliably learn these previously challenging gamma models. Historically, learning such models for long horizons (i.e., when the discount factor gamma approaches one) was considered problematic due to issues with signal degradation and high variance in the learning process. Temporal Difference Flows presents a solution that significantly mitigates these challenges, opening up new avenues for advanced planning, credit assignment, and policy optimization in complex environments.

This work is particularly significant for the ML/systems community as it addresses a fundamental limitation in model-based reinforcement learning: the ability to make accurate and stable long-term predictions. By enabling the effective training of geometrically discounted horizon models, the research paves the way for more sophisticated and robust agents capable of planning over extended horizons. It provides a practical and theoretically sound framework for practitioners to build more powerful world models, potentially transforming how we approach model-based control and decision-making in various AI applications.

Background

▶ Watch: Introduction and problem: compounding errors in world models (0:00)

The foundational problem addressed by Temporal Difference Flows relates to the inherent limitations of traditional world models in reinforcement learning. A standard world model aims to learn the dynamics of an environment, typically represented as the probability of transitioning to a next state given the current state and an action. While indispensable for model-based planning, these models suffer from a critical flaw: compounding errors. As a world model is unrolled over multiple timesteps to predict future trajectories, any small prediction error at one step can accumulate and amplify over time, leading to increasingly inaccurate predictions. This phenomenon severely restricts their applicability for long-term planning, often limiting effective horizons to only five or ten steps in realistic environments.

To circumvent this issue, the research focuses on a different class of world model: geometrically discounted horizon models. Instead of predicting a specific state K steps ahead, these models predict a discounted sum or distribution of future states. This objective is closely related to the successor representation or successor measure in reinforcement learning, which quantifies the expected future occupancy of states under a given policy, geometrically discounted by a factor gamma. Mathematically, for a state s and action a, this model aims to learn P(s_k | s_0, a_0, ..., a_{k-1}) discounted by (1-gamma) gamma^(k-1). An interesting property of this objective is its recursive, Bellman equation-like form. It can be split into two parts: a one-step prediction and a bootstrapped term that relies on predictions from the model itself. Specifically, the target distribution for a future state s' can be expressed as (1-gamma) P(s' | s, a) + gamma ∫ P(s'' | s, a) P(s' | s'', a') ds''. The first term predicts the immediate next state, weighted by (1-gamma), while the second term, weighted by gamma, is a bootstrap from the model's own predictions of future states from the next state.

The classic challenge with these gamma models, particularly as gamma approaches 1 (signifying a desire for longer-term predictions), arises from this bootstrap term. If one tries to learn this model by sampling from it, with probability (1-gamma) one takes the next state, and with probability gamma one samples a state from the model's own prediction. As gamma tends to one, the signal from the one-step transition (the (1-gamma) term) diminishes to zero, leaving the learning process reliant almost entirely on the bootstrapped term. Learning from this bootstrapped term is notoriously difficult because it involves sampling from the model itself, which is often noisy or poorly initialized, leading to a "no signal" problem and high variance. The common wisdom was that these models, while theoretically appealing, were practically intractable for long horizons.

Temporal Difference Flows addresses this intractability by integrating iterative generative models into the learning process. These models, such as flow matching or diffusion models, are designed to transform a simple source distribution (e.g., Gaussian noise, X0) into a complex target distribution (X1) by learning a continuous-time vector field. In flow matching, this is achieved by constructing a conditional probability path, often a linear interpolation, between a source sample X0 and a target sample X1, and then learning the vector field that transforms X0 to X1 along this path. The key insight is that while the ground truth probability path for the entire target distribution might be unknown, if we can sample from the target, we can define a conditional path for individual samples. This framework provides the necessary tools to approach the problem of learning the complex, policy-dependent distributions inherent in geometrically discounted horizon models.

Key Findings

▶ Watch: Introducing geometric horizon models and their Bellman equation (0:50)

The central discovery presented in Temporal Difference Flows is that the seemingly intractable problem of learning geometrically discounted horizon models (gamma models) can be effectively solved by strategically leveraging the structure of iterative generative models, specifically flow matching. This stands in stark contrast to prior common wisdom, which suggested that such models were too difficult to learn for long prediction horizons.

The research's main contributions and findings include:

• Overcoming Intractability with Iterative Generative Models: The paper demonstrates that by framing the learning of gamma models as an objective within the context of flow matching, the challenges associated with the bootstrap term can be managed. The objective boils down to a weighted loss function with two terms, corresponding to the one-step transition and the bootstrap component.
• Introducing "TD-squared" Methods for Variance Reduction: The most significant technical finding is the development of "TD-squared" methods. These methods specifically address the high variance introduced by the bootstrap term, which traditionally requires sampling from the generative model itself. The key innovation is coupling samples: instead of resampling a new initial point (X0) when generating the target (X1) for the bootstrap term, the same initial X0 is reused. This seemingly simple change drastically reduces the variance in the learning gradient, making the training process stable and efficient.
• Quantified Performance Gains: The TD-squared methods show remarkable empirical improvements. In benchmarks, they were shown to be "almost four orders of magnitude better" at estimating the value function for particular policies compared to other methods. This represents a substantial leap in the accuracy and stability of learning these models.
• Improved Prediction Quality: Visualizations in simple environments, such as a 2D maze, clearly illustrate the superiority of TD-squared. While other bootstrapping methods quickly diverge from the ground truth, TD-squared predictions match the ground truth "quite well," providing a clear visual depiction of their effectiveness.
• Theoretical Understanding of Variance Scaling: The paper provides theoretical insights into why the problem is so difficult for large gamma. It shows that the variance in the gradient during the training process scales with gamma^2. This finding validates the intuition that as gamma approaches 1 (longer horizons), the learning problem becomes exponentially harder, and underscores the importance of the variance reduction techniques introduced.
• Enabling Policy-Dependent Planning: The authors highlight that the geometrically discounted distribution of future states is inherently policy-dependent. This property, when effectively modeled, enables a powerful form of planning: by training a gamma model that can predict outcomes under various parameterized policies (e.g., goal-conditioned policies), an agent can query the model to understand the future implications of different actions or policies and select the best one.
• Substantial Gains in Planning Performance: When these learned gamma models are integrated into a policy selection framework, the research demonstrates "pretty substantial gains" in various environments. This translates directly into more effective and robust decision-making for autonomous agents.

In essence, Temporal Difference Flows transforms a theoretically appealing but practically challenging class of models into a viable and high-performing tool for reinforcement learning, primarily through intelligent application of modern generative modeling techniques and crucial variance reduction strategies.

Technical Deep Dive

▶ Watch: The challenge of learning gamma models: signal loss (2:00)

The technical foundation of Temporal Difference Flows rests on framing the learning of geometrically discounted horizon models within the robust framework of flow matching. The primary objective is to learn a vector field that can transform a simple noise distribution into the target geometrically discounted future state distribution. This objective can be decomposed into a weighted loss function with two critical terms.

Let's denote the target geometrically discounted distribution of future states as P_gamma(s' | s, a). The Bellman-like equation for this distribution is key:

P_gamma(s' | s, a) = (1-gamma) P(s' | s, a) + gamma E_{s'' ~ P(s'' | s, a)} [P_gamma(s' | s'', a')]

Here, P(s' | s, a) is the one-step environment dynamics, and P_gamma(s' | s'', a') is the bootstrapped term, representing the future states starting from s'' under a potentially different action a'. The learning objective for the vector field v (which defines the flow) is constructed to minimize the difference between the model's predicted flow and the ground truth flow for both components.

The overall loss L is a weighted sum:

L = (1-gamma) L_one_step + gamma L_bootstrap

1. The One-Step Term (L_one_step):

This term is relatively straightforward. For a given state-action pair (s, a) from the dataset, we have an observed next state s_next. We want our model to predict this immediate transition. In conditional flow matching, this involves:

• Sampling a noise vector X0 (e.g., from a Gaussian distribution).
• Taking the observed s_next as X1.
• Constructing a linear interpolation path between X0 and X1.
• Learning the vector field v such that it correctly transforms X0 towards X1 along this path. This term ensures the model accurately captures the immediate dynamics.

1. The Bootstrap Term (L_bootstrap):

This is where the core challenge and innovation lie. This term requires the model to learn to predict future states from its own predictions. If we have an initial state s_0 and action a_0, the process to generate a target X1 for the bootstrap term would traditionally involve:

• Sampling s_next from the true environment dynamics P(s_next | s_0, a_0).
• Then, recursively sampling from the generative model itself to get a future state s_future starting from s_next (i.e., s_future ~ P_gamma(s_future | s_next, a_next)). This s_future would be our X1.
• Similar to the one-step term, we'd sample a new noise vector X0 and connect it to s_future for conditional flow matching.

The problem with this naive approach is that the process of sampling s_future from the generative model and then sampling a new X0 to connect to it for conditional flow matching introduces a massive amount of variance into the gradient. Each time an update is made, a fresh set of samples is drawn, leading to noisy and unstable learning, especially as gamma approaches 1, making the bootstrap term dominate.

The TD-squared (Coupled Flow Matching) Solution

The breakthrough of Temporal Difference Flows is the introduction of TD-squared methods, which address this variance issue through coupling. The key insight is simple yet profound: instead of resampling X0 for each conditional flow matching path, we reuse the X0 that was initially used to integrate and obtain the bootstrapped X1.

Here's how it works for the bootstrap term with coupling:

1. Start with an initial noise X0_initial.
2. Use the current generative model (defined by the learned vector field) to integrate X0_initial forward to obtain a future state X1_bootstrap. This X1_bootstrap serves as the target for the bootstrap term.
3. Crucially, when applying conditional flow matching to learn the vector field for this (X0_initial, X1_bootstrap) pair, we reuse X0_initial itself as the source point for the conditional path. We do not sample a new X0.

This coupling strategy significantly reduces the variance because the X0 and X1 points are no longer independently sampled. They are intrinsically linked, as X1 is derived directly from X0 through the model's current flow. This practice, common in other flow matching applications for variance reduction, proves to be exceptionally effective here. The authors highlight that the variance of the gradient during training for these coupled methods scales with gamma^2, which provides a theoretical justification for why the problem becomes harder for larger gamma, and why effective variance reduction is paramount.

The choice between conditional and marginal flow matching is also pertinent. While learning a marginal path directly is generally intractable (as it requires knowing the ground truth probability path), the conditional flow matching approach, especially with coupling, makes it feasible. The TD-squared methods essentially enable the model to learn the correct target distribution by minimizing the differences in the vector fields that transform noise to the desired future state distributions for both the one-step and bootstrapped components.

Finally, the talk emphasizes that the geometrically discounted distribution of states is policy-dependent. This means the model learns P_gamma(s' | s, a, pi), where pi is the policy. This policy dependence is a feature, not a bug. It allows the training of a single, generalizable model that can then be queried for different policies. For example, in goal-conditioned policies, the model can be used to check what the future looks like under various potential goal-reaching strategies, enabling effective policy selection and planning. The ability to model this object robustly allows an agent to "check what the future looks like under all of my policies and select the best one."

Experimental Setup & Results

▶ Watch: Leveraging iterative generative models like flow matching (2:40)

The experimental validation of Temporal Difference Flows primarily focused on state-based environments, specifically within the MuJoCo physics simulator. These environments, characterized by their continuous state spaces and complex dynamics, provide a challenging testbed for learning accurate world models. The speaker explicitly noted that the paper's results are based on these state-based settings, where data can be "non-homogeneous... per dimension," which can be "quite hard to make these predictions." While the team has since explored pixel-based environments (finding them "slightly easier to model"), the core contributions in the paper were demonstrated on state vectors.

The primary objective of the experiments was to quantify the effectiveness of the proposed TD-squared methods in learning geometrically discounted horizon models and their subsequent utility in planning. Key aspects evaluated included:

• Value Function Estimation Accuracy: A critical metric for evaluating the quality of learned dynamics models in RL is their ability to estimate value functions. The TD-squared methods demonstrated a dramatic improvement in this regard. The speaker stated that these methods were "almost four orders of magnitude better" at estimating the value function for particular policies compared to other baseline or naive approaches. This translates to a 10,000-fold improvement in precision, highlighting the profound impact of the variance reduction techniques.
• Prediction Quality in Simple Settings: To provide an intuitive understanding of the model's performance, visualizations were presented for simple environments, such as a 2D maze. These visual depictions showed that while other methods, particularly those relying on naive bootstrapping, quickly caused predictions to "go wrong pretty quickly" and diverge significantly from the ground truth, the TD-squared method produced predictions that matched the ground truth "quite well." This qualitative result corroborates the quantitative improvements, demonstrating the model's ability to maintain accuracy over longer horizons.
• Planning Performance: Beyond merely learning the dynamics, the ultimate test for a world model is its utility in decision-making. The geometrically discounted horizon models, once learned using TD-squared, were integrated into a policy selection framework. By allowing the agent to evaluate the future implications of various parameterized policies (e.g., goal-conditioned policies), the system could select the optimal action sequence. The results indicated "pretty substantial gains in these environments when we do this process," suggesting that the improved accuracy and stability of the learned gamma models directly translate into superior planning capabilities.

The experiments did not explicitly detail specific hardware used (e.g., GPU types, number of GPUs) or the exact versions of machine learning frameworks (e.g., PyTorch, TensorFlow). However, the mention of MuJoCo environments implies a standard computational setup for modern reinforcement learning research. The datasets for training these models were derived from interactions within these MuJoCo environments, consisting of "state-action-next-state pair" trajectories used to supervise the one-step dynamics component of the loss function. The absence of specific model sizes or parameter counts in the transcript suggests that the focus was more on the algorithmic innovation and the general applicability of the approach rather than scaling to extremely large models.

In summary, the experimental results unequivocally support the claims that Temporal Difference Flows, through its TD-squared methodology, provides a robust and highly effective solution for learning geometrically discounted horizon models, offering significant improvements in prediction accuracy and subsequent planning performance in state-based control tasks.

Practical Implications

▶ Watch: Weighted loss combining flow matching for gamma models (4:00)

The "Temporal Difference Flows" work carries significant practical implications for a broad spectrum of practitioners in the AI/ML community, including infrastructure teams, model builders, and those involved in deploying intelligent agents.

Firstly, a primary implication is the increased accessibility and usability of geometrically discounted horizon models (gamma models or successor representations). As articulated by the speaker, this class of models was previously "a little bit inaccessible" or challenging to use effectively due to inherent learning difficulties. By providing a stable and reliable training methodology, Temporal Difference Flows transforms these models from theoretical constructs into practical tools. This opens up "huge opportunities" for their application in areas where long-term, discounted predictions are valuable.

For model builders and researchers in reinforcement learning, this means a new, powerful primitive for constructing world models. Instead of relying solely on one-step dynamics or suffering from compounding errors in multi-step unrolls, they can now build models that directly predict discounted future state distributions. This capability is crucial for:

• Advanced Planning Algorithms: Agents can now plan more effectively over longer horizons, as the models provide stable, accurate, and policy-dependent predictions of future states. This facilitates more sophisticated search, tree-based planning, or model predictive control algorithms that were previously hampered by inaccurate long-term forecasts.
• Credit Assignment: The successor representation is intrinsically linked to credit assignment in RL. Robustly learning these models can lead to more efficient and accurate attribution of rewards to past actions, improving policy learning, especially in sparse reward settings.
• Policy Evaluation and Selection: Since the learned gamma models are policy-dependent, they can be used to evaluate the long-term consequences of different policies or actions. This enables agents to perform sophisticated "what-if" analyses, checking "what the future looks like under all of my policies and select the best one" before committing to an action, which is invaluable for robust decision-making.

For infrastructure teams and those concerned with model deployment, the emphasis on "stable and reliable" training is a critical benefit. Unstable training processes lead to unpredictable model quality, requiring extensive hyperparameter tuning and often resulting in models that fail in production. The TD-squared methods, by significantly reducing variance, promise a more robust development cycle, leading to models that are easier to train, validate, and deploy with confidence.

However, certain tradeoffs and limitations should be considered:

• Computational Cost: While the method offers significant stability, training iterative generative models like flow matching can be computationally intensive, especially for complex, high-dimensional state spaces. The integration process of the vector field and the need for potentially large models might require substantial computational resources (e.g., high-end GPUs).
• Generalization to New Domains: The talk primarily focused on state-based MuJoCo environments. While the speaker mentioned that pixel-based environments were explored post-paper and found "slightly easier" for other reasons, the initial challenge for "non-homogeneous data like per dimension" in state-based environments suggests that adapting these methods to radically different data modalities or very high-dimensional observation spaces (like raw video without pre-processing) might still present engineering challenges.
• Gamma Scaling: While TD-squared dramatically reduces variance, the underlying theoretical finding that variance scales with gamma^2 implies that learning truly infinite-horizon predictions (where gamma is extremely close to 1) will always remain challenging, albeit now feasible to a much greater degree. Practitioners need to be mindful of the effective horizon they are targeting and tune gamma accordingly.

In conclusion, Temporal Difference Flows provides a powerful, validated framework that unlocks the potential of geometrically discounted horizon models. Its practical implications range from enabling more sophisticated planning and credit assignment in RL to offering a more stable and reliable approach for training a previously difficult class of world models, ultimately fostering the development of more intelligent and capable AI agents.

Key Takeaways

• Overcoming World Model Limitations: Temporal Difference Flows introduces a robust method to learn geometrically discounted horizon models (gamma models), addressing the compounding error problem that limits traditional world models to short-term predictions.
• Leveraging Iterative Generative Models: The core innovation is the successful integration of flow matching, an iterative generative model technique, to learn the complex, policy-dependent distributions of future states inherent in gamma models.
• TD-squared for Variance Reduction: The "TD-squared" methods significantly reduce the high variance historically associated with learning the bootstrapped component of gamma models. This is achieved through coupling samples by reusing the initial noise vector (X0) when generating and learning from model-predicted future states.
• Dramatic Performance Improvements: TD-squared models demonstrated "almost four orders of magnitude better" accuracy in value function estimation and showed "pretty substantial gains" in planning performance in MuJoCo environments, reliably matching ground truth predictions.
• Enabling Policy-Dependent Planning: The learned gamma models capture policy-dependent future state distributions, allowing agents to effectively query and compare outcomes under different policies for advanced planning and optimal policy selection.
• Stable and Reliable Training: This work transforms a previously inaccessible and unstable class of models into a practical tool for reinforcement learning, offering "stable and reliable" training that opens up new opportunities for planning and credit assignment.

About the Speaker(s)

The work on Temporal Difference Flows was presented by Jesse Farebrother, who conducted this research during an internship at Meta. Jesse Farebrother collaborated with a team of distinguished researchers, including Matteo Pirotta, Andrea Tirinzoni, REMI MUNOS, Alessandro Lazaric, and Ahmed Touati. This collective effort brought together expertise in reinforcement learning, generative models, and theoretical machine learning to tackle a fundamental challenge in model-based control. While specific individual bios beyond Jesse Farebrother's internship at Meta were not detailed in the talk, the collaboration underscores a strong research background in cutting-edge AI/ML methodologies.