Temporal Difference Flows
Authors: Jesse Farebrother
Matteo Pirotta
Andrea Tirinzoni
Remi Munos
Alessandro Lazaric
Ahmed Touati
Venue: ICML 2025 oral (2025)
Keywords: Reinforcement Learning
Geometric Horizon Model
Gamma-Model
Temporal Difference Learning
Successor Measure
Flow Matching
Primary area: reinforcement_learning

Abstract: Predictive models of the future are fundamental for an agent's ability to reason and plan. A common strategy learns a world model and unrolls it step-by-step at inference, where small errors can rapidly compound. Geometric Horizon Models (GHMs) offer a compelling alternative by directly making predictions of future states, avoiding cumulative inference errors. While GHMs can be conveniently learned by a generative analog to temporal difference (TD) learning, existing methods are negatively affected by bootstrapping predictions at train time and struggle to generate high-quality predictions at long horizons. This paper introduces Temporal Difference Flows (TD-Flow), which leverages the structure of a novel Bellman equation on probability paths alongside flow-matching techniques to learn accurate GHMs at over 5x the horizon length of prior methods. Theoretically, we establish a new convergence result and primarily attribute TD-Flow's efficacy to reduced gradient variance during training. We further show that similar arguments can be extended to diffusion-based methods.  Empirically, we validate TD-Flow across a diverse set of domains on both generative metrics and downstream tasks, including policy evaluation. Moreover, integrating TD-Flow with recent behavior foundation models for planning over policies demonstrates substantial performance gains, underscoring its promise for long-horizon decision-making.

Temporal Difference Flows
Jesse Farebrother 1 Matteo Pirotta?
Andrea Tirinzoni Rémi Munos
Alessandro Lazaric 3
Ahmed Touati
Abstract
dictive capability has traditionally been achieved through world models that capture the transition structure of the en- vironment. These models have enabled significant advances across numerous domains from robotics manipulation employing model-predictive control (Sikchi et al., 2021; Hafner et al., 2023; Hansen et al., 2022; 2024), to sample- efficient exploration strategies (Schmidhuber, 1991; Stadie et al., 2016; Pathak et al., 2017), and sophisticated plan- ning algorithms (Silver et al., 2016; 2017; Schrittwieser et al., 2020). However, while world models have demon- strated impressive results, they face fundamental limitations when deployed for long-horizon reasoning. The standard approach of unrolling predictions step-by-step leads to com- pounding errors, as small inaccuracies in each prediction accumulate and propagate forward in time (Talvitie, 2014; Jafferjee et al., 2020; Lambert et al., 2022). This "curse of horizon" presents a significant challenge for applications requiring reliable long-range predictions.
Predictive models of the future are fundamental for an agent's ability to reason and plan. com- mon strategy learns a world model and unrolls it step-by-step at inference, where small errors can rapidly compound. Geometric Horizon Mod- els (GHMs) offer a compelling alternative by di- rectly making predictions of future states, avoid- ing cumulative inference errors. While GHMs can be conveniently learned by a generative ana- log to temporal difference (TD) learning, existing methods negatively affected by bootstrapping predictions at train time and struggle generate high-quality predictions at long horizons. This paper introduces Temporal Difference Flows (TD- Flow), which leverages the structure of a novel Bellman equation on probability paths alongside flow-matching techniques to learn accurate GHMs at over 5x the horizon length of prior methods. Theoretically, we establish new convergence sult and primarily attribute TD-Flow's is efficacy to reduced gradient variance during training. We fur- ther show that similar arguments can be extended to diffusion- based methods. Empirically, we vali- date TD-Flow across diverse set of domains on both generative metries and downstream tasks, in- cluding policy evaluation. Moreover, integrating TD-Flow with recent behavior foundation models for planning over policies demonstrates substan- tial performance gains, underscoring its promise for long-horizon decision-making
An alternative approach is to learn generative model of future states directly, avoiding compounding errors dur- ing inference. These models, usually referred to as Ge- ometric Horizon Models (GHM; Thakoor et al., 2022) or y-models (Janner et al., 2020), are learned by leveraging the temporal difference structure of the successor measure (Blier et al., 2021). However, their reliance bootstrapped predictions during training can lead to instability and grow- ing inaccuracy over long horizons. As a result, current methods struggle to make accurate predictions beyond 20- 50 steps, also limiting their utility for long-term decision- making. In this paper, we show that while state-of-the generative methods like flow matching (Lipman et al., 2023) and denoising diffusion (Ho et al., 2020) cannot be directly applied to learn long-horizon GHMs, their iterative nature can be leveraged to better exploit the temporal difference structure of the problem. This insight yields a new class of methods that provably converges to the successor measure while reducing the variance of their sample-based gradient estimates, enabling stable long-horizon predictions. Empir- ically, our approach produces significantly more accurate GHMs at all horizons, consistently outperforming state-of- the-art algorithms across domains and metrics, including prediction accuracy, value function estimation, and general- ized policy improvement.
1. Introduction
Predictive modeling lies the heart of intelligent decision- making, enabling agents to reason and plan in complex environments. In Reinforcement Learning (RL), this pre-
"Work done at Meta 'MeGill University "Mila Québec AI Institute 'FAIR at Meta. Correspondence Jesse Farebrother <jfarcbro@@ es megill.ca>> Ahmed Touati <atouati @ com>.
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Temporal Difference Flows
2. Background
Geometric Horizon Model A Geometric Horizon Model (GHM; Thakoor al., 2022) or "-model (Janner al., 2020) is generative model of the normalized successor measure. To learn the parametric model m(-.. :0) z m² we can minimize a Monte-Carle cross-entropy objective over source states from the empirical distribution p as,
In the following, we use capital letters to denote random variables, sans-serif fonts for sets, and (A) to denote the space of probability measures over a measurable set A.
Markov Decision Process We consider reward-free dis- counted Markov decision process M = (S, A, 7)) which characterizes the dynamics of a sequential decision- makinn problem. At each step, the agent selects an action a E A in state 8 E S according to its policy π 5 + A. This action influences the transition to the next state s' S. governed by the transition kernel P: S X A (S), which defines probability measure over successor states. The discount factor y E [0, 1) can be interpreted as implying process that either continues with probability or terminates with probability 1 - 7. This interpretation naturally defines a geometric distribution of future states the agent will occupy, where states reached after steps are discounted by y
Exep, X (1).()))))) log m(X | A:
In order to sample from m deploy policy for t Geom(1 7) steps resulting in state X S.. Similar to other Monte Carlo methods in RI., this approach is problem- atic when learning from off-policy data, often resulting in Sighanariance estimators that rely on importance sampling.
A Alternativel, we can leverage the Bellman equation (3) to construct an off-policy iterative method for estimating m Given initial weights 6(9), each iteration updates 0 by minimizing the following iemporal-diiffrence cross-entropy objective over transitions that need not come from policy x,
Successor Measure The normalized successor mea- sure (Dayan, 1993; Blier et al., 2021) of policy describes the discounted distribution of future states visited by start- ing from an initial state-action pair (s, a). For the mea- surable subset X C S the successor measure m² (X 8, a) represents the probability that future states fall within X, geometrically discounted by y according to the time of visi- tation. Formally, it is defined as:
E ((..A)=..-((TTm<*))))) (
(-log log m(X S,): (4)
In the equation above and throughout the paper, we adopt the shorthand m(x)) mm... ; ((x)). To generate samples X ~ (T" m (n)) (-1 S, we first draw a successor state S' P(- S, A): then with probability 1 7, we return S: otherwise, with probability Y, return a bootstrapped sample drawn from (-( S", (S"))
mm m (x)ss))))=
s,a)=
(1)
Several probabilistic models have been applied to this prob- lem, including generative adversarial networks (c.g., Janner et al., 2020; Wiltzer et al., 2024b), normalizing flows (e.g.. Janner et al., 2020), and variational auto-encoders (e.g.. Thakoor et al., 2022: Tomar et al., 2024). We now turn attention to a class of generative models based on the flow- matching framework specifically designed to leverage the underlying structure of the Bellman equation (3), enabling more effective generative models of the successor measure.
(1->)
Σ
* E =s, Ao = a, R),
k==
where Pr(-1 | Ao, π) denotes the probability of state- action sequences (Sk, Ax)x20 generated from (So, Ao) fol- lowing Sp ~ P(- P(-1Sk-1: Ak-1) and A& (St).. The successor measure encapsulates the long-term dynamics of π, enabling value estimation for any reward function r: S + R. Specifically, the value of taking action a e A in state s E is the expected reward under states visited by w amplified by the effective horizon (1 (1-7)--
3. Temporal Difference Flows
Flow Matching (FM; Lipman et al., 2023; 2024; Liu et al., 2023; Albergo & Vanden- 2023) constructs a time- dependent probability path m; : S A + (S) for [0.1] that evolves smoothly from the source distribution mo Po (P)) to the target distribution my R m²" This evolution is governed by a vector field V2 S x x A + S, which dictates the instantaneous movement of samples along my. The relationship between on and the resulting probability path m, established through time-dependent flow v.
Q" (s((s,a) = (1 (1- Ex Ex-m (11 [r(x).
(2)
Moreover, m " is the fixed point of the Bellman operator T : (S))AAA + (S)SxX (Thakoor et al., 2022):
m²( (ssaa) = (T" m²) ) |-| s))a) ==
(3)
s, + (P= s))
The operator pa applied to m mixes the one-step kernel with the successor measure, accounting for transitioning from (8, to new state-action pair (s') and query- ing the successor measure m(- | |s' (s')) thereafter:
A S, defined by the following ODE:
d (x s a) = ve (ve( a) s,a), vo(o s,a) = x de
(,(x)s, a) = x ++/ $ vr ( (1) dr.
(P*m) (dx s,a)= / P(ds' s) m(dx ss.
2

---

Temporal Difference Flows
TD-CFM
Coupled TD-CFM
TD*-CFM
1 (Xol (()())
(Xo 1S (('))
(Xa (((")
Xo
(G))
Mx Xo,
X1
1R6
mig 1 ( 8 ((5'))
mg
my ( 5 ((())
mc
mn (118,x(8))
Figure 1. Visual depiction of TD-Flow variants. Samples are mapped from my to the target distribution mi through the neural ODE vf") Dashed lines depict the neural ODE trajectory; solid lines show the conditional probability path u, (Left) TD-CFM maps Xo to Xi hefore creating separate conditional path between Xú and X1, resulting in erossing paths. (Middle) TD-CFM(C) dircetly couples Xo used to generate X₁ when constructing the conditional probability path. (Right) TD? solves the neural ODE up to time to directly obtain the target velocity es.
We say that is generates my if its flow 8/2 satisfies X, := (e((x | S, A) ~ me(- | S, A) for Xu ~ my- In words, the flow 4/x pushes samples forward through time, ensuring they are distributed according to m, at time t. To learn this transformation, we can minimize the squared I distance between parameterized vector field ie(... ;0) and the true vector field V/ over & ~ (( 4)), yielding the Monte-Carlo Flow Matching (MC-FM) loss (0):
Proposition 1 (Lipman et al. 2024). Given a conditional probability path Pazz and vector field uz2 with their associ- ated marginal counterparts pe(x)) and ve(x), we have
Vo AMC-FAL (0) = Vo (MC.CFM(())
TD-CFM While (MC-CFM; 6) requires direct access to sam- ples from the target distribution m², we can instead learn from an offline dataset p containing only one-step transi- tions (S, S')") through an iterative process similar (4), Starting with initial parameters ((0), at each iteration, minimize the TD-Conditional Flow Matching (TD-CFM) loss &D-CFM an extension of (MC-CFM; 6) that differs only in its sampling procedure:
Ep,,x. [| S, A) - (Xe S, |
where X, (- S,
(MC-FM; 5)
Despite its conceptual simplicity, direct optimization of the flow matching objective above proves challenging due to the inaccessibility of the true probability path my and its associated vector field Uz.
Xo ~ Po
Alternatively, Lipman et al. (2023) shows that we can sidestep this problem entirely by introducing additional con- ditioning information. Instead of directly modeling the
2= = X1 ~ (1-7) 6s + 7
mith (Xo I ((())
(TD-CFM; 7)
In this procedure, with probability 1 - 7, we return the successor state S'. Otherwise, with probability we sample from the neural ordinary differential equation (Chen et al., 2018) 7(x) with corresponding vector field a(()) (X, S', (S")) from Xo ~ Po to produce sample X1 (. (. S', (S')).
probability path m, we can introduce a random variable 7 and define a conditional path on Z as Ps S x - (S) (Lipman et al., 2024; Tong et al., 2024). The conditional ve- locity field wziz= S Z + S that generates Pi1z can now be computed in closed form for many simple choices of Z and Puj2. One such choice is taking Z = XX and performing linear Gaussian interpolation from Xo + X1 resulting in
Coupled TD-CFM Although (TD-CFM: 7) offers princi- pled way of learning the flow from noise to data, an increas- ingly popular strategy to improve flow matching methods is to correlate noise and data whenever a "natural" cou- pling is available (e.g., Liu et al., 2023; Shi et al., 2023; Pooladian et al., 2023: Tong et al.. 2024; De Bortoli et al., 2024). Motivated by this idea, we observe that the process used to generate X2 described above already provides direct coupling between X, and X, We can leverage this coupling by conditioning the probability path Pızz on both endpoints, i.e., Z = (Xo, X)), rather than just condition- ing on = X₁ as in TD-CFM. As illustrated in Figure 1, this coupling helps align X, with the path generated by
(- (-1 (-1X,) = N'( tX,, (1 r) with the correspond- ing vector field given by dejn (x X2) = (X, - z1/(1 t). Armed with the ability to sample from Ptj1 and to compute Wej1. we can directly learn ún by optimizing the Monte-Cario Conditional Flow Matching (MC-CFM) objective Luc-ar (9):
E p.t. x. || (X, A;0) (X, 11
where Z X1 m ( (, X1 p- Z).
(MC-CFM; 6) Remarkably, both (MC-FM; 5) and (MC-CFM; 6) share the same gradient and converge to the same solution.
3

---

Temporal Difference Flows
potentially simplifying the regression problem. This procedure gives rise to the Coupled TD-Conditional Flow Matching (TD-CFM(C)) loss ETDCCMIC) which now extends mm-ers again, differing only in its sampling procedure:
where (e a) - Pip (ar P a)). For the second
a). term, we can leverage the GHM m.)) learned at the previous iteration to construct the marginal,
(n)
0 (z(s. (als', a)) my (1) (c)ss a d') a)
No ~ Po
p. (xls,a)
X1 ~ (1 - 7) dg + 7 8 17 (x.l8"
where p) (Ir s, 5 (x 1 (dd)) | a), and a' od= =s'). This shows that m( plays the role of a condi- tional probability path for the bootstrapped term and vi is its associated conditional vector field. We can then use the equivalence between FM and CFM in Proposition I to replace the marginal probability paths and vector fields in Lemma I with their conditional counterparts to obtain the loss:
Z = (X):
(TD-CFM(C): 8)
(n)
A convenient approach to specifying the conditional path Pazz is to define X, = (Xo, X1) a, + 3o as the affine interpolant between Xn and X1, with the in- terpolation coefficients satisfying the boundary conditions Qg = 31 = 0, a1 = 30 = 1, and monotonicity constraints de 0, - 3. 0, where the over-dot denotes the time derivative. From this definition, the conditional vector field arises as the time derivative of this interpolant defined as (X, Xo. X1) = (A = á + 3. (Al- bergo et al., 2023). A simple choice of the interpolation coefficients that yields a linear (straight-line) conditional path is given by B. = a, = =1=
((0) = 2.
[|, S. 0) - 24 (X, 1 [)]
where 2 = X1 ~ P(: A) X,
71z((112)
2(0) = E 8i [0,(X, S, (s, - 06 1 (R.lS"=(S'))1" ((5))|²],
where Xo ~ Po. S' ~ P(( S. A).
X, = - 7(t)
(Xo (X |S', (5'))
TD²-CFM While (TD-CFM(C); 8) improves upon (TD-CFM; 7) by accounting for the coupling between bootstrapped samples and their generating noise, both methods rely upon fitting an ad-hoc conditional vector field 2z2 that generates the surrogate conditional path Puzz To formulate a more structured approach, we exploit the linearity of the Bellman equation, as detailed in the following result.
1o0 (0) = (1 (1-))((())++) --)0)(++++000
(TD²-CFM: 9)
Since we now bootstrap the previous estimate not only in the sampling process but also in the objective function, we refer to this method CFM). The right panel of Figure I depicts the process of ob- taining the bootstrapped vector field i((*) for TD²-CFM. We provide further implementation details and pseudo-code for all TD-Flow methods in Appendix C.3.1. Next, we extend our TD² result to the class of denoising diffusion models.
TD2-Condiiional Flow Matching (TD²-
Lemma 1. Let Pe be a probability path for P generated by
vector field Th and (m)) be a probability for p*m (n)
generated by Vt (x))) such that po = Po mag For any
((n)
E [0, 1] and () let vi 1. (-),,)) be the solution of
3.1. Extension to Diffusion Models
arg min (1 - 1) x, (1 [I((Xe) - (t s, []
Denoising Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) build diffusion process starting from a data sample Xo ~ = m (- A,) and corrupting it via a stochastic differential equation (SDE),
u: Rd Rp
[1o(2)) - -66" ". (Se s, a1)]
+
Then ii (n+1) induces a probability path m
(a+1) such that
dx,- = f(t) x, dt + dW,,
(10)
(n+1) mo and m (n+1)_ mg
T= (x)
where t E 10, T) for some time horizon T. f.g [0,7] + R is drift and diffusion term, and W, E IRd is a standard Brow- nian motion. The forward process of the linear SDE (10) has an analytic Gaussian kernel (- X₂) = (- of T), where a and o can be computed in closed form. To sample from the target data distribution go. we can solve the reverse SDE (Song & Ermon, 2019) from time T to 0:
This result shows that it is possible to use two independent probability paths for the two terms in the sampling process induced by the Bellman operator. For the first term, we can use a standard CFM approach for Z = X2 with conditional path pan and vector field taln which induces the marginal,
(x x1)) Pan (Ix (d))ss))) Weee f
d.X, = (f(t) (t)(() Vx,log q (X, |S, (s, dt g(()) dw,
T, =
-
(,())s,a)
(11)
Notice here that the minimization is over the space of all functions and not the parameterized vector fields ve(-..;0)
Different flow matching, time is inverted in diffusion mod- els and ranges from 0 to T.
4

---

Temporal Difference Flows
where W, is the reverse-time Brownian motion and qt is the marginal distribution of both the forward (16) and reverse (17) process. To simulate (II), we can train parametrized score function (z s, a; f) to approximate Va, log (xre s,a) using the denoising diffusion / score matching objective (Vincent, 2011) (0):
where B == (1 - P) + 7 m and P,(a)ssa)= == J Pej (xr P(nlls a)dz For any t E [0, 1). the opera- tor Be is a yccontraccionn in I-Wasserstein distance, that is, for any couple of probability paths pt, Qt.
sup W1 Pt) (- 8, ssa)), (57 (- s,))) 8,0
Xp, [|s S,A: - Vx, (X, x. where Xy ~ m. S, 4), X₂ ~ 2:j[( Xo)- (DD; 12)
S Ssup W1 (prl 8, (,))) 2: s,,)).
8,0
Temporal Difference Diffusion Following the blueprint in $3, we define an iterative process starting from g(0) = 5(...: ; ) and minimize at each iteration the Temporal- Difference Denoising Diffusion (TD-DD) loss (o): (0):
Theorem 1 shows that all TD-flow methods fundamentally implement the same update where the probability path at
E [0, 1] is obtained by applying a Bellman- like operator Be to the previous iteration. This operator is a 7-contraction as T directly implying the following result.
X.. ( A, (n0)
Vz log (X xo)l("]
(m) "uzo be the sequence of probabil- ity paths produced by TD-CFM, TD-CFM(C), or TD²-CFM
Corollary 1. Let {m, Inzo
where Xo ~ (T=mb1) (. S. 1), X2 ~ gnjo Xo)
starting from an arbitrary vector field vt Then,
(TD-DD; 13)
(n) Once again, to sample Xo ~ m o/T) (-1 A). we pro-
lim my (n) = =m = =B, 11000
ceed as follows: with probability 1-7, we draw a successor state S ~ P(- A); conversely, with probability Y, we sample from the bootstrapped model by solving the reverse SDE with score function s(n), initiated from Xy. Follow- ing an approach analogous to Lemma 1, we demonstrate in Appendix B that we can employ two distinct diffusion pro- cesses for the two terms involved in the Bellman operator, which consequently leads to the TD -DD objective:
where my is the unique fixed point of Bi, and my = MC where m s.a) = Paul- x1) m (x1 | s,a) is the probability path of the Monte-Carke approach (MC-CFM: 6).
This corollary shows that the fixed point of B, coincides with the probability path generated in Monte-Carle Conditional Flow Matching (MC-CFM; 6), which assumes direct access to samples of m An important subtlety in Theorem is that all algorithms apply the same operator for 2 1, but the result holds for 0 only for TD²-CFM. This means that even starting from the same ((0), the three algorithms may generate different sequences {ml") Jnzo. while still converging 174 In Theorems 5 and 6, we show we can reconcile TD-CFM(C) and TD-CFM with TD² under mild assumption on the form of the initial vector field.
E(0)= E p,t.X,
A: V v 40 (0 & {x,))))))) where Xi ~ 9: 1S),
(0) E S. [| |,, )(-) (Xels' S 7 {""
where XT ur, X. ( ( (S))
L (0) = (1-)))))))))+++ + 7(00)
(TD2-DD; 14)
While Theorem analyzes an idealized version of the al- gorithms, in practice gradients are estimated from samples and the following analysis reveals important differences in their variance. We introduce the (unbiased) sample-based gradients for each of the algorithms,
4. Theoretical Analysis
We now study the learning dynamics of an idealized version of the TD-Flow methods, assuming that the flow-matching loss is minimized exactly at each iteration. Under this as- sumption, each iteration we compute probability path m(.") such that mi (n) = 7"mm("- (1-)) which implies that m(() m. by the contraction property of T The fol-
E [oro-CrM (eaneccr(yymcmeo [m).crm)] Vo ETD-CFM(0)) E IID-CrMcCs (YrD-crMcc))] V (TD-CrMMc)( (0) E 2 (YnD² Ve f
lowing result shows that the overall probability paths mn (n) follow a similar process. Proofs are deferred to Appendix E.
where Y summarizes the random variables involved in the loss definitions in (TD-CFM; 7), (TD-CFM(C): 8), and (TD2-CFM; 9) (see Appendix E.6 for a formal definition of the gradients). We want to compare the total variance of the gradient estimates o² = Tr (Covy [9(Y)]). where Tr denotes the trace.
Theorem 1. For any 2 1, the probability paths generated by TD-CFM, TD-CFM(C), or TD2-CFM satisfy
ma (n+1) (x ()))) (BT ml (n) (x s,a), VL 10.1]
5

---

Temporal Difference Flows
Theorem 2. For any n 2 1 and [o, 11, assume that mik ( (x 8, - I (1 | x1)m (r) (x1 s, a)dx,, then
For single policy, we evaluate how well each method mod- els its i) successor measure and ii) value function. While lower errors in estimating the successor measure are ex- pected to lead to better value estimation, this is not always the case since modeling errors may disproportionally affect states with negligible rewards. Additionally, motivated by our theoretical results, we explore how the probability path's 's design affects our proposed methods' relative performance.
TD-CrM = = 2
IPM +
2² E [To(Cov
Theorem 3. For any n 2 1 and
E [0, assume
(x) (zo, 21 that m s, a. = J Prio. (x 21,
Finally, we examine the scalability of our approach by learn- ing a generative model of the successor measure across a class of parameterized policies derived from the Forward- Backward (FB) representation (Touati & Ollivier, 2021; Touati et al., 2023), a non son-genrrativee model of the successor measure. We conclude by demonstrating how TD2 enables more effective planning for task-relevant policies when per- forming Generalized Policy Improvement (GPI; Barreto et al., 2017), far surpassing the capabilities of FB alone.
s, a.)doodr 9, then we obtain
TD-CEMIC) = =d -TM--FMM + I [Tr(CovzIs.A.x. (Vo. (X,I,,))) " (x1((),
where Z = X1). Furthermore, if we use straight conditional paths, i.e., X,- = 1Xx, + (1 -1-18 t) and the linear interpolant X₁ does not intersect for any 8, a, s', then o TD-CTME.) = =0 TD²-crM'
5.1. Empirical Evaluation of Geometric Horizon Models
In both results, the probability path m; (n) from the previous iteration must be identical for the algorithms being com- pared. The analysis reveals that TD-CFM and TD-CFM(C) suffer from larger variance compared TD²-CFM, which uses the vector field e(x) both to sample X, and as tar- get for the regression problem. This variance gap is "dis- counted by 7², which suggests that the performance of these algorithms would be similar for problems with small horizons but would increase as 7 + 1. The extra variance in both cases stems from samples generated by the algo- rithm (i.e., they do not depend on the transitions available in the dataset). In this sense, we can refer to it compu- tational variance, and in principle, it could be reduced by increasing the number of samples Xo. X1, and X, used in gradient computation. While the variance of TD-CFM and TD-CFM(C) cannot be directly compared, we expect that constructing X, from Xn and X, (instead of X, only) will tend reduce its variance. Specifically, when X, is obtained by linear interpolation between X, and X2, and it does not generate crossing paths, the variance of TD-CFM(C) reduces to the one of TD2-CFM.
Before benchmarking, we must first obtain a policy to eval- uate. We follow the approach taken in Thakoor et al. (2022) and pre-train a set of deterministic policies one for each task - using TD3 (Fujimoto et al., 2018). The final policy obtained from this pre-training phase is now fixed for the remainder of our experiments. GHM training proceeds an off-policy manner where we learn the successor measure of TD3 policy using transition data from the ExoRL dataset (Yarats et al., 2022); specifically, we use dataset of 10M transitions collected by random network distillation policy (Burda et al., 2019). All GHM methods are trained for 3M gradient steps using the Adam optimizer (Loshchilov & Hutter, 2019) with a batch size of 1024 and weight decay of 0.001. We maintain a target network using an exponential moving average of the training parameters with a step size of 0.001. Special care taken to match the capacity of the neural networks between methods with a UNet-style architecture employed for all flow and diffusion methods, while the GAN and VAE baselines use an ML. with residual connections for all their respective networks. Full details for the training methodology, network architecture, and hyper- parameters can be found in Appendix C.
5. Experiments
We implement all conditional flow matching methods (TD- CFM, TD- CC-CFM(C), TD²-CFM) with the Optimal Transport Gaussian conditional path from Lipman et al. (2023). When constructing our bootstrap targets, we sample from the neu- ral ODE using the Midpoint solver with a constant step size of t/ /100 for maximum of 10 steps. For TD²-CFM, we sample ~ (40, 1]); otherwise, we integrate to = 1 and construct X, using the conditional path. For Denoising Dif- fusion methods (TD-DD, TD2-DD), we train DDPM (Ho et al., 2020) by discretizing 3 e (0.1, 20) using T = 000 steps. We construct diffusion bootstrapped targets using
We now present series of experiments assess the efficacy of our TD-based flow and diffusion approaches with base- lines employing Generative Adversarial Networks (Goodfel- low et al., 2014) and 3- Variatoional Auto-Encoders (Higgins et al., 2017). Following the methodology from Touati et al. (2023); Pirotta et al. (2024), we benchmark 22 tasks span- ning 4 domains (Maze, Walker, Cheetah, Quadruped) from the DeepMind Control Suite (Tunyasuvunakool et al., 2020).
#1 (n) (2o, (i(s,a) a) mo ( (m) (mg)a.)) (mr) is the joint "mo1 distribution of (No, X1), ie the endpoints of the ODE.
6

---

Temporal Difference Flows
Scaling Effective Horizon
Method
EMD
Norm NLL
MSE(V)
10
TD-DD TD-VAE
TD-DD
20.22 (0.20)
2.824 (0.195)
454.49 (131.07)
TD3-DD
14.14 (1.08)
0.806 (0.010)
189.15 (23.63)
Value Function MSE
ID-GAN TD-CFM
CHEETAI
TD-CIM
12.26 (0.02)
0.856 (0.021)
228.77 (2.20)
10º
TD-CFM(C)
10.51 (0.0%)
0.447 (0.020)
140.78 (18.72)
TD CFM
10.57 (0.07)
(1.014)
22 (1022))
TD-CFM(C)
GAN
23.97 (0.46)
2163,22 22 (128.00
10
TDD.DD
VAE
83.77 (0.41)
1284.27 CST 622
TD4-CFM
TD-DD
0.149 (0.001)
2.974 (O. 10.10
1215.20 (29
TD--DD
0.027 (0.001)
0.761 (0.082)
11.13 (3.09)
POINTMASS
TD-CFM
0.062 (0.000)
0.554 (0.0333)
355.56 (82.83)
TD-CFM(C)
0.022 (0.002)
-0.99 (0.094)
11.89 (3.16)
TD*-CFM
0.021 (0.000)
-0.833 (0:097)
8.74 (2009)
10"
GAN
0.203 (0.037)
1257.20 (112.800
10
20
50
100
YAE
0.410 (0.038)
1821 889 (09.78)
TD-DD
28.33 (0.33)
1.088 (0:01))
1490.70 (9)
Effective Horizon
TD*-DD
22.64 (2.47)
0.861 (0.0283
159.03 (14.64)
QUADRIPED
TD-CFM
15.73
1.056 (0.0022)
525.06 (28 DO)
TD-CEM(C)
11.38 (0:03)
0.488 (0.005)
155 25 (5 (5,58)
Figure 2. Valuc-Function prediction error as function of the effec- tive horizon (1 -) for y E {0.8, 0.9, 0.95, 0.98, 0.99] on the POINTMASS loop task. TD² methods show impressive robustness to increasingly aag-horizon predictions.
TD*-CPM
14.51 (0.05)
0.379 (6.011)
141.77 (9.10)
GAN
36772.12 (13898.233)
2634.69 (118 38)
YAE
60.27 (0.25)
1156.33 (36.33 (36.52)
TD-DD
20.58 (0.24)
2.649 (0.137)
382 40 (158.63)
TD²-DD
12.09 (0.12)
0.537 (0.060)
39.04 (6.08)
WALKER
TD-CFM
(0.11)
0.713 (0.0283
225, 27 (42.43)
TD- CFM(C)
11.91 (0.02)
0.219 (0.016)
30.77 (3.44)
TD2-CUM
11.92 (0. 109
0.104 (0.001)
28. 35 (0.10)
GAN
24.51 (0.89)
3090.66 (1117.94)
20 steps of the DDIM (Song et al., 2021a) sampler. For TD-DD, we solve to / = 0 and regress towards the noise that re-corrupted our sample. Alternatively, TDD directly re- gresses towards the noise prediction from the target network at a randomly selected noise level. The first bascline we consider a GHM instantiated as a Generative Adversar- ial Network (Goodfellow et al., 2014) similar to the one found in Janner et al. (2020). We follow the best practices from Huang et al. (2024) with the primary modification be- ing a relativistic discriminator Policceur-Martineaul 2019) equipped with a zero-centered gradient penalty on both real and fake samples. For our second baseline, we implement a B- VAE (Higgins et al., 2017) following the practices out- lined in Thakoor et al. (2022).
VAE
111.73((255))
2457.66 61 (16.25)
Table 1. Evaluation results comparing our TD-based methods along with GAN and VAE baselines for single-policy. Results are com- puted over 19 tasks from 4 domains and further averaged across 3 seeds. For each metric we highlight the best performing methods.
Method TD-CFM(C)
EMD
Norm NLL J
MSE(V) 1 310.45 (258.94)
14.08 (12.42)
1.79 (1.98)
TD3-CFM
0.09 (0.09)
-0.01 (0.04)
-3.6 (7.76)
Table 2. Performance difference for TD-CFM(C) and TD²-CFM when employing a curved instead of straight conditional path. Lower is better with negative values indicating a net improvement for using a curved path.
To evaluate the quality of our models, we first generate sam- ples from the ground truth successor measure m" according to the following procedure. We first randomly sample 64 source states So from the initial state distribution and ex- ecute policy æ for 1, 000 steps. Along each trajectory, we resample 2048 states with replacement according to the stop- ping time Geometric(1 7). For the same 64 source states, we generate matching set of 2048 samples from each GHM. Now in possession of these two sets of sam- ples, we evaluate the: 1) log- likelihood of the true samples for models with tractable densities (i.e., diffusion and flow methods); 2) Earth Mover's Distance (EMD; Rubner et al., 2000), which quantifies the minimal transport cost between the two empirical distributions; and 3) mquan- error of a
suggests that our TD-Flow framework should enable more stable training across extended temporal horizons. To vali- date this hypothesis, we train each GHM for 3 seeds on the loop task in the Maze domain while varying the effective horizon (1 - -1 across five values: (5, 10, 20, 50, 100]. Figure 2 illustrates the relationship between value function MSE and the effective horizon. The results demonstrate that - methods maintain consistent performance even as the effective horizon increases, while alternative approaches show significant performance degradation. No- tably, at an effective horizon of 100, TD2-based methods maintain their accuracy and achieve performance improve- ments of nearly four orders of magnitude compared to their naive implementations. These results empirically support for our initial hypothesis, with the stability of TD² methods aligning with our predictions.
estimate of the true value function Q" and the value function derived from GHM samples using (2). Full details can be found in Appendix C.
Having established our training framework, baselines, and evaluation protocol, we proceed to investigate a key predic- tion from our theoretical analysis. Our variance analysis
In the following, we shift our attention to a more in-depth analysis of the largest horizon of 100 (? = 0.99). For each
7

---

Temporal Difference Flows
Planning via Generalized Policy Improvement
FB-GPI
DD-GPI
FM-GPI
FZZ
Coupled
TD2
+36
+36
+36
+36
% Improvement Over FB
+32
+33
4:30
+29
30
+25
20
+16
+12
+10
10
+8
+9
0
-57
-59
-38
-10
Random
Train Distribution
Local Perturbation
Figure 3. Performance improvement over the zero-shot Forward Backward (FB; Touati & Ollivier, 2021) policies when planning with Generalized Policy Improvement (GPI; Barreto et al., 2017). FB-GPI performs GPI over the FB value-function QTT, DP-GPI and FM-GPI perform GPI with the value function implied by the GHM m== for our diffusion- and flow-based methods, respectively. Results are averaged over 22 tasks across 4 domains.
algorithm, train GHM for 3 independent seeds for all domains and tasks. Table 1 reports aggregate performance across our full suite of metrics. For each domain and metric, we highlight results in a 1% range with respect to the best- performing method. The results demonstrate a clear pattern of superior performance for TD2 algorithms: TD?- CFM achieves significant improvements over TD-CFM with 10x reduction in value-function MSE, 1.5x reduction in EMD, and 3 reduction in log- rigllikelihood, averaged across all four domains. In line with our theoretical predictions, the coupled variant of TD-CFM performs comparably to TD CFM, given straight conditional paths. While a comparison between flow matching and diffusion is not at the core of this paper, in our experiments, flow matching generally outper- forms diffusion across all metrics. We posit this primarily due to noise in the diffusion process adversely impacting already noisy prediction problem for large horizons.
7 -x, - Table 2 illus- "rjo, (X, trates the performance difference relative to the straight path results (Table 1) averaged across all domains and tasks. The results strongly support our theoretical prediction: TD²-CFM not only maintained but surprisingly improved performance compared to the linear path. In contrast, TD-CFM(C) showed significant performance degradation, confirming our hypoth- esis about its limitations with non-straight paths.
5.2. Planning via Generalized Policy Improvement
We now turn our attention towards training policy- conditioned GHMs which can be utilized for test-time plan- ning. To accomplish this, first pre-train Forward Back- ward (FB; Touati & Ollivier, 2021; Touati et al., 2023) rep- resentation using the same dataset of 10M transitions as described in $5.1. This pre-training yields a class of w- conditioned policies Tw, where each w E W = (d-1(v represents an embedding of a reward function situated on a d-dimensional hypersphere with radius vd. We then train the GHM - conditioned on the policy by incorporating the embedding w directly into the model's input. All GHM methods are trained for 8M gradient steps, maintaining the same parameters used in $5.1, with the exception of a higher weight decay coefficient of 0.01. For additional insights into the accuracy of the colicy-conditioned GHMs, we direct the reader to Appendix D. Overall, we observed similar trends to those seen in our single-policy experiments.
Given the comparable performance between TD-CFM(C) and TD2 with straight conditional paths, we next ex- amine how these methods behave with alternative path ge- ometries. Our theoretical analysis suggests an important distinction: TD2 should maintain its effectiveness with non-straight paths, while the performance of TD-CFM(C) should degrade. To test this prediction, we maintain the methodology above while replacing conditional path in (TD?-CFM; 9) with the following curved path Pqp (- (X,) with coefficients a = sin (zt) and 8; cos (t. The corresponding conditional vector field is now given by (X, (år - x + X, Ad- ditionally, for TD-CFM(C) we condition the curved path above on Xo and X, resulting in the conditional vector field
W- | X1, )
Given that both FB and w-conditioned GHM models enable estimation of policy's 's value function QTT,, we can utilize this information to perform Generalized Policy Improve- ment (GPI; Barreto et al., 2017) during evaluation. Specifi-
8

---

Temporal Difference Flows
cally, at each time step I. we choose an action az Tw. (se), where We is derived as follows:
2021). Furthermore, recent work on consistency models (Song et al., 2023; Yang al., 2024) and self-distillation (Frans et al., 2025) suggests promising avenues for tackling the computational burden of sampling limitation com- mon to the family of iterative generative models that our approach builds upon.
We E arg max (1 11-7) -- (1 [x()] w>D(W)
(15)
Here D(W) is sampling distribution over W. We consider three such distributions: i) Random: uniform distribution over W; ii) Local Perturbation: we perturb the embedding Wr of the task reward r by the uniform distribution; iii) Train Distribution: we sample w from the training distri- bution used by FB. To approximate (15), we sample 255 embeddings from D(W and explicitly include the task em- bedding w,,, resulting in maximization over 256 policies. To estimate QTT,, we average the reward over 128 states sampled from m Performance is measured by averaging returns over 100 episodes, each lasting 1000 steps.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
References
Albergo, M. S. and Vanden-Eijndent E. Building normaliz- ing flows with stochastic interpolants. In International Conference on Learning Representations, (ICLR), 2023.
Figure illustrates the average percentage of improvement for each algorithm and u-sampling strategy relative to the performance of the FB policy Tw, for the task reward r. We refer to Appendix D for more detailed view of these results. All TD-based GHM approaches lead to signifi- cant improvement over the base FB policy, with TD-CFM(C) and TD²-CFM providing 25 30%+ improvement with all sampling approaches. TDD also leads to significant per- formance gains but is still dominated by the flow matching methods. Notably, FB-based GPI not only fails to improve performance but actually deteriorates it on average with sig- nificant degradation observed in three out of four domains (detailed results available in Appendix D). When compar- ing different distributions D(W). we observe that while FB-GPI's performance fluctuates considerably, GHM meth- ods maintain their robustness across distributions, showing only minor variation. These results underscore the abil- ity of our improved GHMs to make long-term predictions enabling powerful planning capabilities.
Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. Stochastic interpolants: A unifying framework for flows and diffusions. CoRR, abs/2303.08899 2023.
Anderson, B. D. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313- 326, 1982.
Ba, J., Kiros, J., and Hinton, G. E. Layer normalization. CoRR, abs/1607. (.6.6000.660,...0 2016.
Barreto, A., Dabney, W., Munos, R., Hunt, J. J., Schaul, T., Silver, D., and van Hasselt, H. Successor features for transfer in reinforcement learning. In Neural Information Processing Systems (NeurIPS), 2017.
Barreto, A., Hou, S., Borsa, D., Silver, D., and Precup, D. Fast reinforcement learning with generalized policy up- dates. Proceedings of the National Academy of Sciences (PNAS), 1
6. Discussion
2020.
In this paper, we introduced temporal difference flows, a novel generative modeling approach that significantly ad- vances long-horizon predictive models of state. By leverag- ing the successor measure's 's temporal difference structure both in its sampling procedure and learning objective, TD². CFM and TD2-DD effectively address challenges associated with modeling long-range state dynamics. The methods developed in this paper provide a robust theoretical and em- pirical foundation that demonstrates the advantages of our framework across range of tasks, metrics, and domains. We envision numerous exciting applications emerging from this work, particularly around imitation learning (Wu et al., 2025; Jain et al., 2025), planning (Sutton, 1991; Thakoor et al., 2022; Zhu et al., 2024), and ooolpo evaluation (Pre- cup et al., 2000; 2001; Nachum al., 2019; Fujimoto et al.,
Blier, L., Tallec, C., and Ollivier, Y. Learning successor states and goal-dependent values: mathematical view- point. CORR, 0at///2110071222 2021.
Borsa, D.,, Barreto, A., Quan, J., Mankowitz, D. J., van Hasselt, H., Munos, R., Silver, D., and Schaul, T. Uni- versal successor features approximators. In International Conference on Learning Representations (ICLR), 2019.
Burda, Y., Edwards, H., Storkey, A. J., and Klimov, O. Ex- ploration by random network distillation. In International Conference on Learning Representations (ICLR), 2019.
Cetin, E., Touati, A., and Ollivier, Y. Finer behavioral foun- dation models via rego-seee features and advantage weighting. CoRR, abs/2412.04366 2024.

---

Temporal Difference Flows
Chen, T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. Neural ordinary differential equations. In Neural Information Processing Systems (NeurIPS), 2018.
Grathwohl, W., Chen, R. T. Q., Bettencourt, J., Sutskever, I., and Duvenaud, D. FFJORD: free-form continuous dynamics for scalable reversible generative models. In International Conference on Learning Representations (ICIR), 2019.
Dayan, Improving generalization for temporal difference learning: The successor representation. Neural Computa- lion, 1993.
Hafner, D., Pasukonis, J., Ba, J., and Lillicrap, T. P. Mas- tering diverse domains through world models. CoRR, 0bss2301.04104, 2023.
De Bortoli, V.. Korshunova, I., Mnih, A., and Doucet, A. Schrödinger bridge flow for unpaired data translation. In Neural Information Processing Systems (NeurIPS). 2024.
Hansen, N., Su, H., and Wang, X. Temporal difference learning for model predictive control. In International Conference on Machine Learning (ICML), 2022.
Dinh, L., Krueger, D., and Bengio, Y. NICE: Non-linear independent components estimation. In International Conference on Learning Representations (ICLR), Work- shop Track Proceedings, 2015.
Hansen, N., Su, H., and Wang, X. TD-MPC2: Scalable, robust world models for continuous control. In Interna- tional Conference on Learning Representations (ICLR), 2024.
Dinh, L., Sohl-Dickstein, J.. and Bengio, S. Density esti- mation using real nvp. In International Conference on Learning Representations (ICLR), 2017.
Higgins I,, Matthey, L.,, Pal, A., Burgess, C., Glorot, X., Botvinick, M., Mohamed, S., and Lerchner, A. beta- vae: Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations (ICLR), 2017.
Farebrother, J., Greaves, J., Agarwal, R., Le Lan, C., Goroshin, R., Castro, P. S., and Bellemare, M. G. Proto- value networks: Scaling representation learning with aux- iliary tasks. In International Conference on Learning Representations (ICLR), 2023.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion prob- abilistic models. In Neural Information Processing Sys- tems (NeurIPS), 2020.
Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Bois- bunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras,
K., Fournier, N., Gautheron, L., Gayraud, N. T., Janati, H., Rakotomamonjy, A., Redko, I., Rolet, A., Schutz, A., Seguy, V.. Sutherland, D. J., Tavenard, R., Tong, A., and Vayer, T. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021.
Huang, N., Gokaslan, A., Kuleshov, V., and Tompkin, J. The gan dead; long live the gan! a modern gan baseline. In Neural Information Processing Systems (NeurIPS), 2024.
Jafferjee, T., Imani, E., Talvitie, E., White, M., and Bowl- ing, M. Hallucinating value: A pitfall of dyna-style planning with imperfect environment models. CORR, abs/2006.04363 2020.
Frans, K., Hafner, D., Levine, S., and Abbeel, P. One step diffusion via shortcut models. In International Confer- ence on Learning Representations (ICLR), 2025.
Fujimoto, S., van Hoof, H., and Meger, D. Addressing function approximation error in actor-critic methods. In International Conference on Machine Learning (ICML), 2018.
Jain, A. K., Lehnert, L.., Rish, I., and Berseth, G. Maximum state entropy exploration using predecessor and succes- sor representations. In Neural Information Processing Systems (NeurIPS), 2023.
Fujimoto, S., Meger, D., and Precup, D. A deep reinforce- ment learning approach to marginalized importance sam- pling with the successor representation. In International Conference on Machine Learning (ICML), 2021.
Jain, A. K., Wiltzer, H., Farebrother, J., Rish, I., Berseth, G., and Choudhury, S. Non-adversarial inverse reinforcement learning via successor feature matching. In International Conference on Learning Representations (ICLR), 2025.
Ghosh, D., Bhateja, C. A., and Levine, S. Reinforcement learning from passive data via latent intentions. In In- ternational Conference on Machine Learning (ICML., 2023.
Janner, M., Mordatch, I., and Levine, S. Gamma-models:: Generative temporal difference learning for infinite- horizon prediction. In Neural Information Processing Systems (NeurIPS), 2020.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farles D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. In Neural Information Processing Systems (NeurIPS), 2014.
Jolicoeur Martinur,, A. The relativistic discriminator: a key element missing from standard gan. International Conference on Learning Representations (ICLR), 2019.
10

---

Temporal Difference Flows
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference Learning Representations (ICLR), 2015.
successor representation. Journal of Machine Learning
Research (JMLR), 24:80:1-80:69 2023.
Misra, D. Mish: A self regularized non-monotonic neural activation function. CoRR, htwp/1988.86111.. 2019.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. In International Conference on Learning Repre- sentations (ICLR), 2014.
Nachum, O., Chow, Y., Dai, B., and Li, L. Dualdice: Dehavior-agnostie estimation of discounted stationary distribution corrections. Neural Information Process- ing Systems (NeurIPS), 2019.
Lambert, N., Pister, K., and Calandra, R. Investigating com- pounding prediction errors learned dynamics models. CoRR, 0tss/203..966377 2022.
Park, S., Kreiman, T., and Levine, S. Foundation policies with hilbert representations. In International Conference on Machine Learning (ICML), 2024.
Le Lan, C., Tu, S., Oberman, A., Agarwal, R., and Belle- mare, M. G. On the generalization of representations in reinforcement learning. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2022.
Pathak, D., Agrawal, P., Efros, A. A., and Darrell, T. Curiosity-driven exploration by self-supervised predic- tion. International Conference on Machine Learning (ICML), 2017.
Le Lan, C., Greaves, J., Farebrother, J., Rowland, M., Pe- dregosa, F., Agarwal, R., and Bellemare, M. G. A novel stochastic gradient descent algorithm for learning princi- pal subspaces. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2023a.
Perez, E., Strub, F., de Vries, H., Dumoulin, V., and Courville, A. C. Film: Visual reasoning with gen- eral conditioning layer. In AAAI Conference on Artificial Intelligence, 2018.
Le Lan, C., Tu, S., Rowland, M., Harutyunyan, A., Agarwal, R., Bellemare, M. G., and Dabney, W. Bootstrapped rep- resentations in reinforcement learning. In International Conference on Machine Learning (ICML), 2023b.
Pirotta, M., Tirinzoni, A., Touati, A., Lazaric, A., and OI- livier, Y. Fast imitation via behavior foundation models. In International Conference on Learning Representations (ICLR), 2024.
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In International Conference Learning Representations (ICLR), 2023.
Pooladian, A.-A., Ben-Hamu, H., Domingo-Enrich, C., Amos, B.. Lipman, Y., and Chen, R. T. Q. Multisample flow matching: Straightening flows with minibatch cou- plings. In International Conference on Machine Learning (ICML), 2023.
Lipman, Y., Havasi, M., Holderrieth, P., Shaul, N., Le, M., Karrer, B., Chen, R. T. Q., Lopez-Paz, D., Ben- Ben-Hamu, H., and Gat, I. Flow matching guide and code. CoRR, 0bs/2412.062664 2024.
Precup, D.,, Sutton, R. S., and Singh, S. Eligibility traces for off off-policy policy evaluation. In International Conference on Machine Learning (ICML), 2000.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. In International Conference on Learning Representations (ICLR), 2023.
Precup, D., Sutton, R. S., and Dasgupta, S. Off-policy tem- poral difference learning with function approximation. In International Conference on Machine Learning (ICML), 2001.
Loshchilov, I. and Hutter, F. Decoupled weight decay reg- ularization. In International Conference on Learning Representations (ICLR), 2019.
Rezende, D. and Mohamed, S. Variational inference with normalizing flows. In International Conference on Ma- chine Learning (ICML), 2015.
Machado, M. C., Rosenbaum, C., Guo, X., Liu, M., Tesauro, G., and Campbell, M. Eigenoption discovery through the deep successor representation. In International Confer- ence on Learning Representations (ICLR), 2018
Ronneberger, O., Fischer, P., and Brox, T. U-net: Convolu- tional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Inter- vention (MICCAI), volume 9351, pp. 234-241, 2015.
Machado, M. C., Bellemare, M. G., and Bowling, M. Count- based exploration with the successor representation. In AAAI Conference on Artificial Intelligence, 2020.
Rubner, Y., Tomasi, C., and Guibas, L. J. The earth mover's distance as a metric for image retrieval. International Journal of Computer Vision, 40(2):99-121 2000.
Machado, M. C., Barreto, A., Precup, D., and Bowling, M. Temporal abstraction in reinforcement learning with the

---

Temporal Difference Flows
Schmidhuber, A possibility for implementing curiosity and boredom in model-building neural controllers. In International Conference on Simulation of Adaptive Be- havior, 1991.
Stadie, B. C., Levine, S., and Abbeel, P. Incentivizing ex- ploration in reinforcement learning with deep predictive models. In International Conference Learning Repre- sentations (ICLR), 2016.
Schramm, L. and Boularias, A. Bellman diffusion models. CoRR, 2ttp//407..121.... 2024.
Sutton, R. Dyna, integrated architecture for learning, planning, and reacting. ACM SIGART, 2(4): 1991.
Schritwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., Lillicrap, T., and Silver, D. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604 609, 2020.
Talvitie, E. Model regularization for stable sample rollouts. In Conference on Uncertainty in Artificial Intelligence (UAI), 2014.
Thakoor, S., Rowland, M., Borsa, D., Dabney, W., Munos, R., and Barreto, A. Generalised policy improvement with geometric policy composition. In International Confer- ence on Machine Learning (ICML), 2022.
Shi, Y. De Bortoli, V. Camphell, A., and Doucet, A. Diffu- sion schrödinger bridge matching. In Neural Information Processing Systems (NeurIPS), 2023.
Sikchi, H., Zhou, W., and Held, D. Learning off-policy with online planning. In Conference on Robot Learning (CoRL), 2021.
Tirinzoni, A., Touati, A., Farebrother, J., Guzek, M., Kan- ervisto, A., Xu, Y., Lazaric, A., and Pirotta, M. Zero-shot whole-body humanoid control via behavioral foundation models. International Conference on Learning Repre- sentations (ICLR), 2025.
Silver, D., Huang, A., Maddison, J. Guez, A., Sifre, L., van den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M., Dieleman, S., Grewe, D., Nham, J., Kalchbrenner, N., Sutskever, I., Lillicrap, T. P., Leach, M., Kavukcuoglu, K., Graepel, T., and Has- sabis, D. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587):484-489, 2016.
Tomar, M., Estruch P., Bachman, P., Lamb, A., Langford, J., Taylor, M. E., and Levine, Video occu- pancy models. CoRR, abs/2407.09533 2024.
Tong, A., Fatras, K., Malkin, N., Huguet, G., Zhang, Y., Beotor-Brooks, J., Wolf, G., and Bengio, Y. Improving and generalizing flow-based generative models with mini- batch optimal transport. In Transactions Machine Learning Research (TMLR), 2024.
Silver, D., Schritwieser, J.. Simonyan, K., Antonoglou, I, Huang, A., Guez, A., Hubert, T., Baker, L., Lai, M., Bolton, A., Chen, Y., Lillicrap, T., Hui, F., Sifre, L.., van den Driessche, G., Graepel, T., and Hassabis, D. Mastering the game of go without human knowledge. Nature, 550(7676) 2017.
Touati, A. and Ollivier, Y. Learning one representation optimize all rewards. In Neural Information Processing Systems (NeurIPS), 2021.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International Conference on Machine Learning (ICML), 2015.
Touati, A., Rapin, J., and Ollivier, Y. Does zero-shot rein- forcement learning exist? In International Conference on Learning Representations (ICLR), 2023.
Tunyasuvunakool, S., Muldal, A., Doron, Y., Liu, S., Bohez, S., Merel, J., Erez, T., Lillicrap, T., Heess, N., and Tassa, Y. con_tontol: Software and tasks for continuous control.
Song. J., Meng, C., and Ermon, S. Denoising diffusion implicit models. In International Conference on Learning Representations (ICLR), 202
Software Impacts, 6: 100022, 2020.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. In Neural Information Processing Systems (NeurIPS), 2019.
van den Oord, A., Vinyals, O., and Kavukcuoglu, K. Neural discrete representation learning. In Neural Information Processing Systems (NeurIPS), 2017.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score- based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR), 2021b.
Vincent, P. A connection between score matching and de- noising autoencoders. Neural Computation, 23(7):1661- 1674, 2011.
Wiltzer, H., Farebrother, J., Gretton, A., and Rowland, M. Foundations of multivariate distributional reinforcement learning. In Neural Information Processing Systems (NeurIPS), 2024a.
Song, Y., Dhariwal, P., Chen, M., and Sutskever, Consis- tency models. In International Conference on Machine Learning (ICML), 2023.
12

---

Temporal Difference Flows
Wiltzer, H., Farebrother, J., Gretton, A., Tang, Y., Barreto, A., Dabney, W., Bellemare, M. G., and Rowland, M. A distributional analogue the successor representation. In International Conference on Machine Learning (ICML), 2024b.
Wu, R., Chen, Y., Swamy, G., Brantley, K., and Sun, W. Diffusing states and matching scores: A new framework for imitation learning. In International Conference on Learning Reprexentations (ICLR), 2025.
Yang, L., Zhang. Z., Zhang, Z., Liu, X., Xu, M., Zhang, W., Meng, C., Ermon, S., and Cui, B. Consistency flow matching: Defining straight flows with velocity consis- tency. CoRR, abs/2407. 023999, 2024.
Yarats, D., Brandfonbrener, D., Liu, H., Laskin, M., Abbeel, P., Lazaric, A., and Pinto, L. Don't 't change the algorithm, change the data: Exploratory data for offline reinforce- ment learning. CoRR, abs/2201.13425, 2022.
Zhang, Chen, X., Zhao, L., Xiong, W., Qin, T., and Liu, T.-Y. Distributional reinforcement learning for multi- dimensional reward functions. In Neural Information Processing Systems (NeurIPS), 2021.
Zhu, C.. Wang, X., Han, T., Du S. S., and Gupta, A. Dis- tributional successor features enable zero-shot policy op- timization. In Neural Information Processing Systems (NeurIPS), 2024.
13

---

Temporal Difference Flows
Appendices
A. Related Work
The Successor Representation (Dayan, 1993) was originally proposed for tabular MDPs and was later generalized to continuous state spaces with the Successor Measure (Blier al., 2021). Successor Features (Barreto et al., 2017; 2020) extends these ideas by instead modeling the evolution of multi-dimennional features assuming rewards decompose linearly over these features. Prior works have leveraged these methods for zero-shot policy evaluation (Dayan, 1993; Barreto et al., 2017; Wiltzer et al., 2024b), zero-shot policy optimization (Borsa et al., 2019; Touati & Ollivier, 2021; Touati et al., 2023; Park et al., 2024; Zhu et al., 2024; Cetin et al., 2024; Tirinzoni et al., 2025), imitation learning (Pirotta et al., 2024; Jain et al., 2025), exploration (Machado et al., 2020; Jain et al., 2023), representation learning (Le Lan et al., 2022; 2023a;b; Farebrother et al., 2023; Ghosh et al., 2023), and building temporal abstractions (Machado et al., 2018; 2023).
Janner et al. (2020) originally proposed a method learn a generative model of the successor measure with modeling techniques spanning from Generative Adversarial Networks (Goodfellow et al., 2014) to Normalizing Flows (Dinh et al., 2015; Rezende & Mohamed, 2015) like RealNVP (Dinh al., 2017). Followup work (e.g., Thakoor et al., 2022; Tomar et al., 2024) explored other generative modeling techniques including various types of auto-encoders (e.g., Higgins et al., 2017: van den Oord et al., 2017). Also of note is recent work learning generative models of multi-dimensional cumulants including features (Wiltzer et al., 2024a; Zhu et al., 2024) and multi-variate reward functions (Zhang et al., 2021). Prior work by Wiltzer et al. (2024b) sought to deal with the instability of long-horizon predictions in GHMs by employing an n-step mixture distribution where they sample at time l along the trajectory. Without resorting to importance sampling this approach is limited to the on-policy setting, Finally, most closely related to our work is that of Schramm & Boularias (2024) who provide a preliminary and limited derivation of what we term TDD In contrast, our work not only rigorously formalizes and significantly extends these ideas but also integrates them into the more general flow-matching framework (Lipman et al., 2023; 2024), additionally incorporating extensions to score-matching (Song et al., 202 bbb) and diffusion (Sohl-Dickstein et al., 2015; Ho et al., 2020). Moreover, we conduct an extensive empirical analysis, demonstrating the efficacy of our approach - an aspect notably absent from Schramm & Boularias (2024).
Geometric(1 -) and bootstrap if > n; otherwise returning the state
B. Extension to Score Matching and Diffusion Models
This section extends our framework to score matching and denoising diffusion models. We leverage the unification of these methods under stochastic differential equations (Song et al., 2021b) introducing an analogous class of Temporal Difference Diffusion methods.
B.1. Background
Both score-based generative modeling (Song & Ermon, 2019) and diffusion probabilistic modeling (Sohl-Dickstein et al., 2015; Ho et al., 2020) can be unified under the framework of stochastic differential equations (SDE) introduced in Song et al. (2021b). Unlike in flow-matching, time is inverted in diffusion models and ranges from time 0 to T. Given the data distribution go and prior simple distribution qr (the "noise" distribution), we construct diffusion process (Xi)ee(or)) such that X, go and XT dr- This diffusion can be modeled as the solution an Ito SDE:
dx,- = f() X, df + g(t) dW, | Xo ~ qo
(16)
where W, is standard Brownian motion and f [0 + Rd is scalar function called the drift coefficient, and g: [0, T] + R is scalar function known as diffusion coefficient.
Generating samples from X, ~ go consists in sampling X² ~ qr and reversing the forward-SDD. process in (16). A known result from Anderson (1982) states that the reverse of a diffusion process is also a diffusion process, running backward in time and given by the reverse-time SDE:
dx,- = ((t) Xx - g(t)² Vx, log ( + g(t) dW. Xy ~ qr:
(17)
where Wr e a Brownian motion when time flows backwards from to 0, dt is an infinitesimal negative timestep and q
14

---

Temporal Difference Flows
the marginal distribution of X, Therefore, we learn the score of the marginal distribution V from To by simulating the reverse diffusion process (17).
log qe (z), we sample
To estimate V, log (z), we can train time-dependent score-based model 8(..-; ;0) [0. R R via the denoising diffusion/ / score matching objective (Vincent, 201 I; Song & Ermon, 2019):
en = [.-2(( X ax [X, Vx, log Tejo xo)]l²]
(18)
For EnD to tractable, we need to know the conditional probability Φrjo- Usually, specific choices of the drift and diffusion coefficients ft and 9t are used such that Q:jo is always Gaussian distribution (- ou of), where the mean a and variance of can be computed in closed-form. The global minimizer of (0) denoted by 27 (1) is equal to the score function V log (2), thanks to the following proposition:
Proposition 2 (Vincent 20111 1). Let (e(x)) = f (1o) Jaqo (3r (r)) dro then we have:
V hp (0) Vo L (x 0) Vx, log y (X))))
(19)
B.2. Temporal Difference Diffusion
To learn a predictive model of m² using diffusion from an offline dataset, we follow a similar approach to what we presented in $3 and we define an iterative process starting from initial weights g(0) and at each iteration minimizing the Temporal-Dificen Difference Denoising Diffusion (TD-DD) loss:
LDD (0) Ep,,,Xo.X, | (. V log log x1)
(TD-DD; 20)
(n) ( (: (: X, ~ q0ol Xo).
where, Xo ~
(n) In order to sample No (TT (-1 8, a), with probability 1 7, we return the successor state S' ~ P(- S, Otherwise, with probability y we solve the following ceverse-time SDE from Xy using the score s(!)
dx, = (f()) X₁ g(t)2s/ (X, S, dt g())dW.
(21)
Minimizing (0) leads to score function s s(n+t)(s s, a) generating a marginal probability q; (x+1) that approximates T=00 (n)) at O.
Following the TD2 blueprint, we can further exploit the structure of the target bootstrapped distribution to design an improved diffusion process that converts Gaussian noise T T=) First, show below that the mixture of diffusion process is also diffusion process with modified drift and diffusion functions.
Lemma 2. Consider two diffusion processes with drift functions ] and J, sharing the same diffusion coefficient g:
dX, = F((x,) dt g() dw
dX, f((x)) + g(t) dw.
Let qi and qi be their marginal distribution, then the diffusion process corresponding to the mixture marginal distribution St (1 7)% + 7ht is:
dX;- = (1 (1-7) + qi.e (X.) g())dW dw.
(1 (1-7)4 7) +r
Proof. The marginal probabilities p and f are characterized by the Fokker-Planck equations:
0p. div
9
div(pr (FrJe) +
Ajr
Dt
2
Up. div(pefi)
g
2
Api
Ot
15

---

Temporal Difference Flows
where div is the divergence operator and A div V is the Laplace operator. Therefore,
Pp. at
0p: 0L
0pt δt
(1 7)
+
- div (pefe) + Li
+ dh App
Apt div(p.fi)
2
2
-div ( - 7) + +Dpef.
zi A A(1 -)) + 7p))
div pt (1 (1-7) + +nPrfii
Apr.
(1 7 Pt + 7itt
I and the diffusion coefficient Ir satisfy the Fokker-Planck equation with the probability path pts
The drift (1-7)pp tip: and therefore their associated diffusion process generate Pt-
Lemma 2 can be easily extended the case of a continuous mixture of diffusion processes.
This result shows that it is possible to two independent diffusion processes for the two terms in the sampling process induced by the Bellman operator. For the first, we can use the standard noising diffusion process:
(.(x s, / 0efo(1 P(ds's s,
where sample X2 9tjo(- | s') by simulating simple forward diffusion process (16). For the second term, can leverage the GHM mi (n) at the previous iteration to construct the process,
1))
(x s)
=f
f m (x x()) P(ds' s,a),
where (x (x s')a') is the marginal probability of the reverse SDE induced by the score s(,),
dx, = (( - g(1)² (") (X, s.))) + g()).
(m) (xr 8', (ss)) weighted by P(s' 8,0), Additionally, Lt nx (x 8, a), as continuous mixture of diffusion's marginals m²
can be generated by the diffusion process,
dx, x
s, a)) dt + gt) dw,, y, where
6e s,a) = f P(ds' | (5, qe (n) (o 8 (b') s (m)) (nr ((,))()))
P(ds' s, qi (x ss,)(s'))
Given these two diffusion processes, the target probability de (x+1) = (1 7) + 7 (n) can be generated by the following reverse SDE,
ax,- = (f(()x, o(t)²
s6 1 (X, (X, 8, a)) at dW,
where of (n+1) (xr s, (1-2) v. log +20 (())(())))
(x s.)). Therefore, we can leam B.(... ;0) to approximate 8, (x+1) by
(1-1)) (.)))
minimizing the loss,
{()) (1 E.) (1(1.AA)
[|5(X,
S,A; Vx, (x, (S, A)ll [|3(X; As 5. (X₂) S, As..)]']]
(22)
7E p.t. )( (-15.4)
16

---

Temporal Difference Flows
We simplify the first term via Proposition (since (crls, a) [ P(ds'| (s, a)), hence we have
Va B (.
[|x 8, a, 0) - Vx. x. log 2e (X, S,, A)]"
E (.....xanna(((5)
[|s(Xx 5, A; Vx, log (go |
Moreover, using a similar argument for equivalence between the gradient of marginal and conditional Hlw-matching objectives, we can show that
||5(X| S. 31"h (X, S,A)]l²]
Ve E a
Vo X... [|5(X, S,A;;) (x, S, A)1l8-
This leads us to the final 2DD loss function,
-p (8) (1 - x. (X, Vz VologrrooX 13) ||5(Xx11 A;,) -x, 5 (
(23)
+ yE (
17

---

Temporal Difference Flows
C. Experimental Details
C.1. Evaluation
Evaluating a GHM can be challenging, TD-based losses employing bootstrapping do not provide a good signal the quality of the learned model. Instead, we opt measure 1) the likelihood of trajectory coming from the true discounted occupancy of a given policy, 2) the Earth Mover's Distance (EMD; Rubner et al., 2000) be- tween samples from the true occupancy and our GHM which provides an estimate of the distance between these two probability distributions, and 3) the value-function approximation error. In all cases, to obtain samples from the true discounted occupaney,
Table 3. Evaluation hyype--paraeeter for both single and multi nulti-policy experiments.
Evaluation
Hyperparameter
Value
Number of states so
64
Number of m-samples per state
2048
EMD
Number of episodes per state
Episode length
1000
Number of state 80
64
collect trajectories ) from policy and subsequently re- sample states according to /. ~ Geometic(1 - y) for particular discount factor E [0, 1). Armed with samples from m we compute the aforementioned metrics follow- ing the procedures stated below along with the parameter values outlined in Table 3.
{(0,, $1..... 81
Number of GHM-samples per state
2048
MSE(V)
Number of episodes per state
Episode length
1000
Number of samples
256
GPI
Number of GHM samples
128
Number of FB inference samples
250, 000
Normalized Negative Ligeliiilloodd. To compute the log- legllike of our flow matching and diffusion meth- ods, we take advantage of the following change in vari- ables formula (Dinh et al., 2015; Rezende & Mohamed, 2015; Chen et al., 2018),
log (mix s, 8)) = log (xx)) + L-"
Ô llo (m( s, az 0)) dt,
Ōr,
where y is the probability density function of a standard Gaussian distribution, which acts as the prior on Zy- The change log density over time can be written as the following differential equation called the instantaneous change of variables formula (Chen al., 2018, Theorem 1).
log (m(x1 0)) = - -Tr (0 18,0, az;
(
Ox
Or
We can now compute the log-likelihood for sample X m. (s|ss, a) by integrating the total change in log-density backward in time from = X to obtain zo which has tractable likelihood. In practice, we solve the following coupled initial value problem using numerical integration (Grathwohl et al., 2019),
1."
(xe s. a; @ielzz s.a;0) diee
30 mm1 |s,a; - log (
-
dl
Tr
(24)
where
21
X
log m(x a. log m(x1 a, (;:0)
For all experiments we report the negative log-likelihood normalized by the dimension of the observation space.
Earth Mover's Distance We compute the Earth Mover's Distance (EMD; Rubner et al., 2000), also known as the Wasserstein-1 1 distance, between
2048 samples from the ground truth distribution X ~ m (- S As) and our learned GHM X m Sk, Ax; ()) for a set of randomly sampled state-action pairs ((Sx, As) As)] Intuitively, the EMD quantifies the minimum cost required transform one distribution into another, where the cost defined in terms of the Euclidean distance between states (1) X(), Formally, we have,
EMD((X(1, x(m) (.)()()))) (8(1)
{(m) min
EEj
x
x(1)
EEE
i.j
A=
18

---

Temporal Difference Flows
where is transport plan such that Eiy specifies the proportion of mass moved from X; to Xy. We report the average EMD 64 source states using the Python Optimal Transport (Flamary al., 2021) library.
Value Function Mean Square Error (MSE(V)). We compute the mean square error between Monte-Carlo estimation VC of the value function (s) and the estimation VaHM obtained using the learned model. We obtain VM by collecting
trajectory {(80, . }T)} from policy x and computing the discounted sum of rewards. We generate a single trajectory since both the policy and the environment are deterministic. The GHM estimate is given by (2), i.c.,
GHM (s) = - 7) "IE Xxm((11 r(x)]-
Then, MSE(V Vinx) = Es (Veina VisC (So))) We average our results over 64 initial states So sampled from the initial state distribution v.
over 100 episodes, each lasting 1, 0000 steps, for every task. For the Norward-Backwartt representation (Touati & Ollivier, 2021), directly follow the policy Tw, (thus = Tw, (s,)) where = R)sp B(s) R is the zero-shot policy embedding inferred using 250, 000 transitions labeled with the task reward function r. Given that FB provides direct way of estimating the value function of a policy (i.e., Q," (*((s,a) = F(s, a, ") 2,), we can do planning in the policy embedding space by solving the following problem:
Planning with GPI. We evaluate planning performance by computing the average
FB-GPL
arg max F(s (st)), w)² Wr. w-D(W)
This optimization problem requires no generation except sampling from D(W). We approximate the max using 255 samples from D(W) and additionally incorporating Wr to ultimately maximize over 256 policies. On the other hand, for GHM-GPI, we solve following optimization problem,
W
GHM-GM
E arg max (1 71
[x())
w>W((W)
which requires generating samples from m "w. In our experiments we generate 128 samples from m²"
C.2. Environments
Experiments in this paper were conducted with a subset of domains from the DeepMind Control Suite (Tunyasuvunakool et al., 2020) highlighted in Figure 4.
Figure 4. A visual depiction of each domain used in our experiments from the Deep-indd Control Suite (Tunyasuvunakool al., 2020).
From to right: MAZE, CHEETAH, QUADRUPED, WALKER.
19

---

Temporal Difference Flows
C.3. Geometric Horizon Models
This section describes each class of generative model used for our empirical experiments.
C.3.1. FL.w MATCHING
Table 4, Summary of how different TD-flow algorithms generate the target probability path and vector field. The neural ode th is defined by the vector field e. computed at iteration n.
Algorithm 1 Template for TD-Flow algorithms
I: Inputs: offline dataset D. policy æ, batch size n, Polyak coefficient Ç. weight decay A, randomly initialized weights 6, discount factor y, leamnngg rate 1). one-step conditional path Pij1 and conditional vector-field Mtj1. bootstrap path Pt: and vector-field 84
P
2: for 1,... do
3: Sample mini-batch {[Sk, Ax, S from
Xo ma
4: for k L..... K do
TD-CEM
5: Sample 4 ~ [4[0, 1])
X1 = (No A(;0)
Man(Xe X1)
6: Sample Xx ~ P S%
7:
(0) Rt, (Xpl Sa, As; 0) 11 |SDIl
X, pup( |xi)
8:
Sample Xx ~ pip (:1 St (St); 0)
9:
(0) = p., (X&l Sk. Agi 0) Ve, (8x 11, 0S))
TD-CFM(C)
Xo mo
10:
end for
X1 = (X A1;0)
"yoo Xi)
11:
# Compute loss
12:
{(0) = : Ex (1 - --) (0) (0)
X. Pijo, Xx)
13:
#Perform gradient step
14:
0 v (c(0) +
TD CCFM
Xo mo
15:
# #Update parameters of target vector field
oxX, (',
16:
1 F- (0 + (1 <) <)0
X1 = (X ,,,4000)
17:
end for
To discuss the TD-Flow methods introduced herein, we first unify the loss function through defining a general template for the loss as,
{(0) = - 7) P.a.X. Pan (15)) [[ [|,,(X, S,A;0) - (X, 5711"]]
+ 7B +BB p [l S, A; - (X, 2)
(n) and vector field 0, as illustrated We can now recover each algorithm by a specific choice of the target probability path p. in Figure 4. Based on this unified structure, we present pscudo-code for the TD flow methods in Figure 1. In practice, instead of proceeding through full iterations, we use standard mini-batch gradient updates with a target network 0 updated as a moving average of 0.
When employing the conditional probability path Pej1 and vector field Maj1 use the standard Gaussian linear interpolation defined Pep X1) = (- tX1, (1 - 1)) hence X, tX (1 Xo Рμj. resulting in dep (X2 Xi) (X1 (1) - t) (Lipman et al., 2023). The source distribution for all experiments is mo)-) = N- 0, I). To sample
from the Neural ODE we use the Midpoint method with a constant step size of dt = t 1/10 for a total of 10 steps. We found
both coupled and TD2 methods do not require many solver steps and hypothesize this is due to the reduction in transport cost as analyzed Appendix E.
For all flow and diffusion-basee methods, we employ a U-Net-style architecture (Ronneberger al., 2015) that has hierarchical skip connections throughout MLP. We embed the timestep t by first increasing its dimensionality with sinusoidal embedding before transforming it through a two-layer MIP with mish activations (Misra, 2019). We further process additional conditioning information, such as the state-action pair and BawwaddBBakwwer additional two-layer MLP, whose result then gets concatenated with our time embedding. Finally, the network integrates all prior conditioning information through FiLM modulation (Perez et al., 2018) that replaces the learned affine transformation for layer normalization (Ba al., 2016).
embedding 2 through an
20

---

Temporal Difference Flows
C.3.2. DENOISING DIFFUSION
We train Denoising Diffusion Probabilistic Model (DDPM; Ho al., 2020) using the same architecture as our flow matching model above, with the output being interpreted prediction of the noise seed co that began the diffusion process. We discretize the diffusion process using 000 steps with Amin 0.1 and _max 20. We employ the DDIM sampler(Song et al., 2021a) with 50 sampling steps for both training and evaluation.
For evaluating our DDPM model, we compute exact log-likelihoods using the instantaneous change of variables formula (Chen et al., 2018) along with the probability flow ODE from Song et al. (2021b). That is, we solve the initial value problem in (24) using the vector field.
1 (8min t (Smax Pminn) (x),
VT- cx s, a.:)))
(e s, a;
=
We now outline the losses for each of the TD-DPM experiments in the paper:
TD-DD To train our vanilla Diffusion GHM we employ the standard DDPM-style objective, that is, optimize the following loss:
[||-- [ + vI - āge 0.
p. VG E
(25)
(T- m m (115,))
where 0 are the target parameters and are the standard diffusion coefficients as seen Ho et al. (2020).
TDD As outlined $3.1 we can split our DDPM loss into two terms, one that will use standard DDPM training on one-step transitions and the second term that will regress our target networks noise prediction. This materializes as,
ē(0) E,,,,,
[leelvaroo + v1 dre AS =11²
where Xo P()
((0) = E P. | S,A:0) c,-)) (Xt S {(S)))1?)
where /i (.8'",((''))
(0)= (1-7)((()) 2((
(26)
C.3.3. GENERATIVE ADVERSARIAL NETWORK
We implement a modern Generative Adversarial Network (GAN; Goodfellow al., 2014) baseline based on the recom- mendations in Huang al. (2024). Specifically, we train relativistic GAN (Jolicoeur-Martineall 2019) resulting in the following loss,
lGAN(OG, 0p) = Xp.X..x, [f (D(G(X, |S S,A:0c);Op) - D(x, | [,A;
where Xo ~ N/ 0 /) X1 (T-m(") (: S,A)
We take f(x) log (1 + exp (-z)) to be the log-sigmoid function (Jolicoeur-Martinen 2019) and further add the following zero-centered gradient penalties on the discriminator,
R1 (ep) = E p. (45)
[l] V| {[])
R2 0p) = ( S.A.a)
x D(X ()
The penalty Rj penalizes the gradient norm of the discriminator on "real data" sampled from our current iterate (m), whereas R2 penalizes the gradient norm on "fake data" generated directly from the current generator. We experimented with different coefficients and schedules for these gradient penalties and settled on linear decay schedule from 0.05 - 0.005
21

---

Temporal Difference Flows
throughout training. Furthermore, is common practice, impose schedule the second moment EMA coefficient /3 of Adam (Kingma & Ba, 2015) to increase from 0.9 -> 0.99 throughout training.
The generator and discriminator architecture in our GAN is implemented as Residual MLP with leaky ReLU activations with the same FiLM-style conditioning (Perez et al., 2018) our flow and diffusion models. The input to our generator random noise sampled from an isotropic Gaussian with dimensionality equal to the number of state dimensions in the environment.
C.3.4. VARIATIONAL. AUTO-HNCODER
We implement 3- V-rriaationll Auto-Encoder (Kingma & Welling, 2014; Higgins et al., 2017) following the best practices outlined in Thakoor et al. (2022). That is, train our VAE minimize the following loss,
(vAE(Og,0p) 0D) = Ex,x, Exo (1S. log Peb (X1 S,), (loo
where X1 (TT (- A).
We employ similar architecture to GAN-GHM and use residual MLP for the encoder and decoder. We use an isotropic Gaussian latent space with the number of latents equal to the number of state dimensions in the environment. We also swept over 3 E {0.1,0.4 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0} on the MAZE domain and chose B 0.5 for the rest of our experiments. Overall, we found the E-VAE-based GHM to be very unstable and likely requires very careful fine-tuning of 33 to get adequate performance at long-horizons.
22

---

Temporal Difference Flows
C.4. Hyperparameters
We report the hyper-parameters for training the GHM models used in the single and multi-policy experiments. Table shows the parameters for Flow Matching and Denoising Diffusion. We also report the Forward-Backward representation (Touati & Ollivier, 2021) utilized in the multi-policy GHM experiments in Table 8.
for pre- training the
Table Flow Matching and Denoising Diffusion hyper-parametes parameters used for the single and multi-poliey experiments across tasks and domains. We highlight any differences depending on the training context.
Hy Hyperparameter
Single Policy
Multi-Policy
ODE Solver
Midpoint
Midpoint
Flow Matching
(Lipman et al., 2023)
ODE dt (train)
0.1
0.1
ODE di (eval)
0.1
0.05 (0.1 for GPT)
Aaia
0.1
0.1
Ann
20
Diffusion (DDPM) (Ho et al., 2020)
Discretization Steps
1,000
1,000
(Song al., 2021a)
(Song al., 2021a)
SDE Solver Steps (train)
20
20
SDE Solver Steps (eval)
20
t-Positional Embedding Dim.
256
f- Positional Embedding MI.P
(256, 256)
(256, 256)
Network (U-Net)
(Ronneberger al., 2015)
Hidden Activation
mish (Misra, 2019)
mish (Misra, 2019)
Blocks per Stage
Block Dimensions
5112,
(1024, 1024, 1024)
Encoder Input
8,a
6,4,
Conditional Encoder
Encoder MLP
(512, 512, 512)
(1024, 1024, 1024)
Enzoder Activation
mish (Misra, 2019)
mish (Misra, 2019)
AdamW A
0.9
0.9
A2
0.999
0.999
Optimizer (Adam (Loshchilov Hutter, 2019)
AdamW
10-4
10-4
Learning Rate
10
10
Weight Decay
10
10
Gradient Steps
3M
8M
Common
Batch Size
1024
1024
Target Network EMA
10 3
23

---

Temporal Difference Flows
Table 6, 3-VAB (Higgins et al., 2017) htyper-parameter for single policy experiments across tasks and domains.
Table 7. GAN hyper-parameters for single policy experiments
across tasks and domains.
Hyperparameter
Value
Hyperparameler
Value
Grad. Penaliy Coef
Linear(0.05 -+ 0.005) N)
3-VAE (Higgins al., 2017)
Latent Prior
Uolicocu--Martinsee 2019)
Latent Prior
[S]
Latent Dimension
Encoder
Residual MLP
Generator
Rosidual MLP
Discriminator
Residual MLP
Decoder
Residual MLP mish (Misra, 2019)
Network
Hidden Activation
Leaky ReLU
Network
Hidden Activation
Blocks per Stage
Blocks per Stage
Block Dimensions
(512, 512,511)
Block Dimensions
(512, 512)
Encoder Input
8,
Encoder Input
5,0
Cunditional Enculer
Encoder MLP
(512, 512)
Conditional Encoder
Encoder MLP
(512, 512, 512)
Encoder Activation
Leaky ReLU
Encoder Activation
mish (Misra, 2019)
Adam
AdamW &1
0.9
AdamW 32
Lincar(0.9 - 0.99)
AdamW
0.999
Optimizer (Adam')
AdamW
10-4
Optimizer (AdamW)
(Loshchilov Hutter, 2019)
(Loshchilov & Hutter, 2019)
AdamW:
10-4
Learnnng
10-4
Learning Rate
10-4
Weight Decay
10-2
Weight Decay
10
Gradient Steps
SM
Gradient Steps
3M
Common
Batch
Common
Batch Size
1024
Target Network EMA
10-3
Target Network EMA
10
Table 8. Forward Backward Representation hyper-parameters, highlight any deviations.
We largely reuse the hyper- seper-parameters from Pirotta et al. (2024) and
Hy Hyperparameter
Walker
Cheetah
Quadruped
Maze
Embedding Dimension d
100
50
100 5²
Embedding Prior
S4
g
5
Forward Backward
(Touati & Ollivier, 2021)
Embedding Prior Goal Prob.
1/2
B Normalization
Eg
8g
Orthonormal Loss Coeff.
Target Policy Noise
((0.0.2)
W 0.2)
(.,
0.2))
Policy (TD3)
Target Policy Clipping
0.3
0.3
0.3
0.3
(Fujimoto et al., 2018)
Policy Update Frequency
Learning Rate (F, B)
(10-1, 10-1)
(10-1 10
(10-5, 10
(10-1,
10-6)
Leaming Rate (x)
10
10
10
10
Optimizer (Adam) (Kingma & Ba, 2015)
Adam A
0.9
0.9
0.9
0.9
Adam Ba
0.999
0.999
0.999
0.999
Adam
10-6
10-8
10-8
10-6
Batch Size
2048
1024
2048
1024
Gradient Steps
3M
3M
3M
5M
Common
Discount Factor
0.98
0.98
0.98
0.99
Target EMA
0.99
0.99
0.99
Reward Inference Samples
250. (00)
250, 000
250, 000
250, 000
24

---

Temporal Difference Flows
D. Additional Experimental Results
In this section, we report additional results about the experiments.
Single policy. We report metries averaged over tasks using curved conditional path in Table 12. We also report the performance per task in Table 13. Table 11 shows the performance of the single-policy experiments ($5.1 in the main paper) expanded for each task. While the performance of TD-based methods is reasonably stable across tasks, VAE and GAN have large variance across tasks. For example, the EMD of GAN diverges in 2 tasks out of 4 in QUADRUPED.
Multiple policies and planning, We report aggregate performance across our full suite of evaluation metrics for the multi-policy experiments in Table 14. We also report per-task metrics in Table 16. We can notice that TD²-DD achieves quite high EMD compared to TD-DD while achieving a better MSE(V). By further inspecting the generated samples (see Figure 5), we found that TD-DD tends to generate highly concentrated samples, while TD3-DD is more diffuse. However, the samples generated by TD-DD appear to be better at visual inspection. This may explain the discrepancy between the two metrics. Finally, we report aggregate planning performance in Table 15 and per-task results in Table 17.
Comparison with planning with one-step world model we include in Table 9 results for a Model Predictive Path Integral (MPPI) controller with learned dynamics model. We train a similar capacity dynamics model to that of TD²-CFM before evaluating MPPI with a finite horizon of 32 for locomotion tasks and 128 for maze, where at cach step we sample 256 action candidates and perform 10 optimization rounds with 64 elites (top-k actions) per round. The results show that GPI with TD²-CFM 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.
Impact of number of ODE integration steps we report in Table 10 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 show that TD²-CFM remains robust even at coarse discretizations of the ODE with as little as 5 integration steps, while we observe with a predictable degradation when the number of steps is too small,
Table Comparison with planning with one-step world models
Table 10. Ablation the number of ODE integration steps
Domain
FB
TD*-CFM-GPI
MPPI
ODE Steps
NLL
EMD J.
MSE(VF) 4
Cheetah
479.35 (11.56)
693.63 (5.50)
541.22 (5.28)
-0.48 (0.21)
0.076 (0.003)
379.52 (81.75)
Pointmass
472.45 (14.40)
800.99 (8.56)
286.43 (54.95)
-2.75 (0.15)
0.036 (0.000)
23.82 (2.05)
Quadruped
627.28 (1.98)
695.73 (2.07)
156.80 (122.89)
10
-2.85 (0.17)
0.025 (0.001)
7.71 (2.75)
Walker
526.66 (5.94)
627.6 63 (7.97)
658.15 15 (21-46)
20
-2.99 (0.04)
0.0218 (0.001)
4,40 (0.82)
25

---

Temporal Difference Flows
Table II. Per task results for the single policy experiments.
Walker
Quadruped
Task
Method
EMD:
NLL1
MSE(V)
Method
FMD
NLL.1
MSSRV).
TD DD
20.06 (0.27)
2.713 (8.00)
120.21 (52.9
27.89 (000)
1,80 Q11129)
1778.75 (6LL2
TD--DD
(09))
0.543 164
21.02 (15.50
11" D3
(8.72
0,006 00.008
12.88 (2.07)
TD-CIM
12.46
0.088 (0.000)
128.06
15,08
1,088
523,10
RUN
TD CPM(C)
10.90 (0.03)
(0000)
9.53 (L97)
PUMP
TD cPstlc)
14.12 (0.00)
0,018 00,002)
10.30 C 12)
TO--CEM
(111))
(7.00
TD7-CrM
(0000
0,426
12,80
GAN
(1
827 79 (130.39)
GAN
00.54
18.23 (0.30)
55-6.311 (ANEEST)
1939.62
VAE
114.55 (2.61)
646.96 (21.57)
YD 20
26.00 (1.00)
1,05 00.0ET
G18 (300.60
DD
21,55
2.754 (1000))
1812.90 (2016.0)
TDP-po
20.79 (3.00)
Tn'-no
13.57
0.561 (1-01-)
21.81 (9.50)
TD OPM
15.74 (0.00)
1,051 (0,020)
170.66 (10.012
TD-CIM
(1.26
0.808 (00021)
J7S 32 (1500S
TD-CFM0C)
14.62 (0:11)
0.457 00.000
CFM(C)
(2.94(0000)
(0.000)
(L.06
10 CPS1
14.75 (0.00)
0,335 00.004
18.36 (2.02)
TDP-CPM
(x10)
200
7,14 719
DAM
19.21 (0.19)
195.11 (144.29)
GAN
(1.96
2948 74 (3541.00)
VAR
00,56
425.00 02 (10 483
VAE
(6.90)
437.70
Th-30
28.57 (0.50)
1.832 (0.031))
2053 27 (1707 (X))
19,82
(0.0)
16.11
Tp--po
20,81 (1.81)
0.967 01/410
20.09 (HL,00)
TD' DD
12,15
0.487 (0.040)
(1.75
15.00 10.00220:10
1:003 (1.050)
505.51 (S9 (77)
TD-CFM
13:30 (3.30
0 (0:011)
32.9 $2.95 303
STAND
TD-CFMIC))
18.01 (1000)
0,683 010009
12.96 (1
STAND
TD-CIM(E)
12:25 (0:11)
12.76
TDP-DIM
14.07 20.199
0.383 0117
7.97 (0.01)
CPM
12.27 (0.14)
0.126 (0.0299
14.96 (10.20)
20 (51009/81)
3611 LD (2090.
GAN
22.08
504L 85 (104
VAR
50. (0.00000
55A.51 (101
(1.89
3963.70
26.83 (144)
1.004 Quarrty
1000.02
02 (112.01)
Tp" pg
21.96 (1.710
0,8LD (1,0400
570.78
(55,995
TD
1D-00
21.29 (0.45)
(1.89
0.558 (100S))
2.505
(0.072)
88 (24.9g)
WALK
TD-CFMEC)
14,50 (1.01)
(1100)
1,108 01u0c0
694 00.000
SP3.02 (24.553
900.76
12.31
(1005-))
TDP-CFM
(1:319
0,561
$28.06
WALK
TD-CFM(C)
11.55 (3.02)
0.225 (0.050)
(14.00
GAN
55777 67
(18193 15
S100. 15 (14 (3)
TD CFM
11.54 (0.27)
0.118 (1 020
79.41 (16.24)
00.57 (UU54)
1597.62 (1002200
VAE
(1.65)
4888. 07 (78.5)
Pointmass Maze
Cheetah
Taok
Method
MSR(V)
Task
Method
EMD
NLL
MSE(V)
TD-D0
011899
J.40Z 01 2005
20.31 (0:313
2,00
60 001880
TD"-00
TD"-nD
14.44
0,758 (0.122)
172.00 635-9
TD-CPM
0.0R1
(5000)
0.768
GFT. 88
TD-TTM
11.50 11.00-00033
0,888 (11001)
211.22 (35-20)
1.009
TD-CRMIC)
0,005
(1.0095
-0 WOS (1.00O
10.91 (2:35)
PLY
TD-CPMCC)
10.35 (1.00)
0.853 (0.024)
124.08 (17.59)
TDP-CRM
0.000
-0654 OLIMO
1.75
GAN
0.225
(60054)
25R0.20 (RILING
TD' -CFM
10.67 (0.01)
67.76 (21000)
YAH
0.426 43/845
23.53 (9:51)
TD-D0
(1-099
11 808 61 (588)
300 SO Do
VAE
80.00 (1.0)
3339.00 (44.80)
T0--0n
(0000)
TD-D0
16:67 (0:07)
2.047
(1) 1810
1043.27 (3TA 677)
TD-CTM
0.00
(10003
0.520
224.15 DIS 10y
TO"
12,50
0.894 (0L02T)
963,04 (97.08)
BOTZOM LEFT
REACH
TD-CPMIC)
TDP-CTM
0.001
0.000 (30034)
(140095
-0 739
0.964
1677
16.58 112
10.44 (T.00)
TD-CEM(C)
TD-TTM
10.91 17
0.827 (1L01T)
0.542
316 (SRi0
410.09 (77.00)
GAN
0,302 (6.350)
L 50 (2220))
TDa CPM
10.11 (0.14)
0.542 (0L005)
370.6B (112 299
(195.70)
GAN
20,00 (1.95)
3761.79 (15s S1)
TD-DD
0.174
3.370
61 1577
750 79 (18 Ri
84.63
918.52
TD--00
(00000)
0.640
TD-D0
70 (0-01)
2:00
(0.3010
06. (1..5
BOTTON RIGH?
REACH
TD-CPM
0.06 (30044)
0.340
QUARSS)
106.07 DS Thy
B0AN (2.69)
TD-CPM
TDa po
16.91 (4.04)
12.21 (0.05)
0.813 (0.025)
0,872 (0.022)
54.95 (11.00
96,33 (75.4
TDP-CTM
0.000
(000143
0.656 DLEGZO
RUN
TD-CPM(C)
10.44 (0:02)
0.434 (0.015)
24.32 (5 (0)
GAN
416. 75 (14 120
ID
10.53 (0.05)
0,412
27,00 (5:44)
TD-DD
0.505 6046119
1. (1:0013
3 A0G 6148585
55G 9% (TA
OSR
GAN
26.48 (201)
183.47 (12.300)
To--00
(00000)
0.863 eclassp
83.91 (0:573
999
TD-CPM
0.0G5 (0004)
0.464
472.64 (335.83)
TD-DD
21.47 (0,52)
(0.8019
20.26 (6:00)
TOP LEFT
KRACH
YD?-CRM
GAN
(0.0099 0.005
(5020)
UUST
L850, 40 (110 FO)
25 AS (5 4D
BACKWARD
RUN
TD-CMM(C)
TDa pD
13.04 (1.22)
18.58
11.02 (1.003
0 (0U0L
054
14,87 (2.94)
8.71 (1 (1.05
604999
(198 1ay
TDP-CRM
(1:00)
11-00
0.141 (5:008
2.984
61 2415
262 56 (3:01)
GAN
24.77 (1.41)
270.11 (1
6,58
VAF
89.91
(0:303
754.7 (22.91)
700 RIGHT
BRACH
TD-CRMIC)
TD CPM
0.002
0.100 00-015
(1004)
0,30L
-U
771
287 57 (67.189
6.18
TD--np
21.57 (0.84)
12.85
2,700 (0.151)
(2700000117)
646.06 (55/300
238.01 (1LLTT)
TBP-CRM
00018 0310
3:31 (9.00
TD-CFM
12.27 113
0.800
(0.031)
377 45700 01)
GAN
0,215 (11044)
1043.01 O57 M0
WALK
TD-CPM(C)
10:24 (0.17)
0.354
(0.021)
176. 90 (25.50)
0.453 (0.306)
1229.57 (56:68)
TDP-CFM
10.18 (0.000
0.530 (1 0713
220 (21
24.39(1.11)
3520 (1000.70
VAF
84.30 41
21.58 17.55 (133 013
T0-D0
21.03 (0.20)
2.854
(0.054)
TD2
14,61 22.49
0,771
00.0119
160,12 (12.20)
TD-CFM
12.59 (4:
0.857
(0.013)
70 (63.50)
BACKWARD
WALK
10.58 (0.01)
0.412
(0.0289
99:30 (4
TD2-CPM
10.36 (7)
0.381 (0.01
100,00 (10 45
GAN
(0.84)
3434.
47)
VAB
89.73 (1.61)
465.72 (15.00)
26

---

Temporal Difference Flows
Table 12. Results averaged over tasks for the single policy experiments with a curved conditional path.
Domain
Method
EMD-
NLL
MSE(V)-
TD-CPM
13.91 (0.75)
1.351 (0.017)
177.89 (60.53)
CHEETAH
CD-CPM(C)
25.86 (18.9))
1.295 (0.067)
180.21 (17.60)
TD*-CFM
10.79 (0.03)
0.412 (0.014)
121.67 (5.08)
TD-CFM
(0.003)
(0.061)
758.16((108. (103.54)
POINTMASS MAZE
TD-CPM(C)
0.080 (0.008)
4.340 (0.456)
079.29 11).11)
TD--CFM
0,021 (0.000)
00806 (0.017)
9.22 (1.40)
TD-CFM
15.63 (0.09)
1.478 (0.058g
273.68 (31,07)
QUADRUPED
TD-CFM(C)
34.00 (6:96)
0.933 (0.006)
522.28 (150.42)
TD-CIMM
14.56-(0.02)
0.327 (0.01:4)
142.18 (9.38)
TD-CIM
13.10
1.147 (0.012)
608.47 (124.62)
WALKER
TD-CFM(C)
33.20 (4.66)
1.039 (0.052)
189.66 (102.80)
TD TD²-CFM
(2.00(0.05)
0.099 (0.005)
27.56 (0.53)
Table 13. Per task results for single policy experiments with a curved conditional path.
Walker
Quadruped
Task
Method
EMD
NLL-
MSE(V) 1
Task
Method
EMD
NLL
MSE(V) -
TD CFM
12.39 (0.17)
1.218 (0.107)
326.14 (56.12)
TD-CFM
15:31 (0.17)
1.460 (0.188)
115.99 (138.56) 69)
RUN
TD-CTM(C)
23.80 (1.91)
0.923 (0.191)
69.39 (8.81)
JUMP
TD-CFM(C)
39.28 (8.90)
0.980 (0.062)
686.51 (314.19)
TDMM
10.69 (0.06)
-0.040 (0.008)
11.69 (1:01)
TD²-CFM
14.36 (0.07)
0.358 (0.010)
10.84 (3.05)
TD CFM
14.08 (0.12)
1.410 (0.189)
806.83 (278.52)
TD-CFM
15.61 (0.16)
1.450 (0.060)
104.52 (33.53)
SPIN
TD-CTM(C)
47.39 (3:14)
1.801 (0.186)
401.22 (321.52)
RUN
TD-CFM(C)
40.27 (7.50)
0.898 (0.040)
240.50 (58.83)
TDMMCFM
13.37 (0:11)
0.198 (0.008)
7.65 (2,29)
TD3-CFM
14.73 (0.06)
0.288 (0.015)
21.13 (3.52)
TD-CFM
13.24 (0.23)
0.896 (0.0533
274.60 (121:20)
TD-CFM
15.24 (0.11)
1.515 (0.21.5)
173.07 (34.00)
STAND
TD-CTM(C)
36.32 (13.80)
0.625 (0.053)
159.74 (31.53)
CFM(C)
22.77 (6.86)
0.924 (0:053)
275.03 (240.91)
TDP-CFM
12.50 (0.17)
0.119 (0.008)
9.42 (1.95)
TD2-CFM
14.17 (0.09)
0.342 (0.019)
7.05 (1.80)
TD-CFM
12.69 (0.20)
1.067 (0:015)
936.30 (86.71)
TD-CFM
16.37 (0.10)
1.486 (0.022)
701.13 (83.58)
WALK
TD-CFM(C)
25.29 (7.62)
0.808 (0:049)
128.28 (65.70)
WALK
TD-CTM(C)
33.68 (1.69)
0.917 (0.036)
887. 18 (120.82)
TDP-CFM
11.42 (0.20)
0.119 (0.026)
81.47 (3.53)
TD2-CFM
14.99 (0.08)
0.318 (0.018)
529.71 (15.40)
Pointmass Maze
Cheetah
Task
Method
EMD
NLL
MSS(V) L
Task
Method
EMD J
NLL
MSE(V) 1
TD-CTM
(1.15))
1.465
(1)71)
1888.54 (136/00
TD-CFM
12.92 (1.25)
1.324 (0.042)
342.71 (129.009
LOOP
TD-CPM(C)
0.132 (0.031)
5,191 (1328)
1354.09 (102.55)
PLIP
TD-CFM(C)
22.90 (15.00)
1.364 (0.108)
140.32 (42.14)
TD3.CFM
0.020 (1,000)
0.708 (0.01))
2.31 (0.50)
TD' CCFM
10.89 (0.08)
0.433 (0.012)
74.31 (6.30)
TD-CFM
9,096 (0.012)
1.91 (0.142)
626.74
(118.04)
REACH
TD-CEM
14.32 (1.08)
1.316 (0.190)
576.31 (169.57)
BOTTOM LEFT
TD-CPM(C)
TD2-CBM
0.078 (0.001)
0.022
0.883 (0,05))
(0.5710
$20.02 (52 88)
10.55
BACKWARD
ELIP
TD-CFM(C)
25. (85
1.427 (0.027)
388.45 (87.18)
TD-CFM
0.097 (0.001)
1,296 (o. (0.2200)
290.21 (29.94)
TD3-CFM
10.48 (0.25)
0.538 (0.034)
283.84 (10.81)
REACH
TD-CPM(C)
0.109 (0.000)
5.310 (0.502)
400.20 28 (10.79)
TD-CFM
14.00 (0.77)
1.300 (0.043)
114.51 (1-11)
TD'-CFM
0.019
-0.833 (0.040)
2.64 (0.30)
RUN
TD-CPM(C)
17.42 (5.78)
1.423 (0.001)
37.23 (8:74)
TD-CFM
0.070 (0.000)
0,894 (0.209
500.63 (142.18)
TD3-CFM
10.85 (0.08)
0.405 (0.010)
32.58 (8.42)
TOP LEFT
REACH
TD-CPM(C)
TD3-CFM
0,048 (0.000)
0.025 (1,002)
2.821 (0:208)
0.738 (0.011)
76. (20.06)
26.56 (1.90)
KUN
TD-CFM(C)
TD-CEM
38.06 (28.90)
14.50 (0.31)
1 1,283 (0.1100
1.439 (0.100)
101.24 (140.88)
100.32 (5.35)
BACKWARD
TD-CTM
0.083 (0.001)
1.035 (0.138)
482 68 (128.45)
CFM
11.06 (0.05)
0.199 (0.007)
12.32 (2.34)
TOP RIGHT
REACH
TD-CPM(C)
0.060 (0.001)
4.138 (0.305)
737.30 (23.75)
TD-CFM
13.60 (0.71)
1. (0.041)
1040.43
(6)
TD²-CFM
0.019 (0.001>
-08666 (0.0
4.02 (1.75)
WALK
TD-CFM(C)
21.01 (16.43)
1. 1.006 (0.028)
343.71 (66.91)
TD*-CFM
10: (0.04)
0.323 (0.010)
213.87 (23.09)
TD-CEM
13.83 (0.89)
1. 1,396 (0.0333
681.05 (21.3)
WALK
TD-CPM(C)
30. 29 (22 589
1.178 (0.206)
124.20 (17.61)
BACKWARD
TD*-CFM
11.00 (0.11)
0.372 (0.015)
113.09 (22.41)
27

---

Temporal Difference Flows
Table 14, Per domain results for the quantitative multi- policy experiments.
Table 15. Per domain results for the multi-policy experiments evaluating planning performance with generalized policy improvement.
Domain
Method
EMD
MSE(V)
Z- Distribation D
TD-
(7.7
1442
(0.02)
534.22 (107.80)
Random
Local Perturbation
Train Distribution
10"-DD
(7.49
0.771 (0.000)
253,8 (21-42)
FB
470.35 (34.30)
CHEETAR
10-E7M
(200
1044
10.010
826.64 (58.01)
975.52
(5.99)
(111
0.581
(19.81)
1D-00
GPI
074.05
(3.88)
G84 (L1.8)
620 72 (14.29)
7D
11.06
200.3
CHERTAI
TD² DD
GMI
062.17 (0.94)
680.22 (6.983
678.98 (1,97)
VD-00
0.162 (0.00)
2048
(0.009)
662.36 (76.86)
TD'-p0
0.349 (0/07))
0.966 (0.027)
312.98 (66.40)
TD-CIM(C)
GPI
681.52 (6.19)
700497 (6.37)
(67.81 (1.10)
POINTMASS
(0,003)
0.771 (0.021)
TD CPM
GPI
682.1 (6.41)
692 72 (7.00
CORL,33 (5.50)
0.174
(14.40
TD²-CFM
0.060 (0.00)
0.043 (0.002)
160.74 (65.66)
FB
-0.84 (7.70)
240. 54 (23.09)
-17.74
TD-DD
20.21 (1779
1.403 (1-079
400 58 1190-17)
TD-B0
GM
569.05 (37-09
580 90 (37 (N0
537 60 (17.50)
TD--DD
(922
0.901 (0001)
415.29 (101 80)
POINTMASS
ya'-00
763.95 (18.02)
N05.72 (2.05)
788.8T (17.15)
QUADRUPED
(1006
0.960
625. (29122)
671.50 (22799)
099.70
TD-CFM(C)
(015)
0.328 (1.010)
176 (33.59
TD-CFMIC)
GPI
800.57 (1.00
812.44 389
808 00 (1-77)
TDF-CFM
14.74
0.340 (1.010
178.95 (10.43)
TDY -CTM
790.34
813.90
500.00 (1.50)
TD-DD
21.49 21.49 (3U4)
1.441 (0.0009
571.72 (196.7))
TD--DD
0.668 (1.0011)
(17.89)
FB
GPI
671.95 (0.58)
674.00 08 (0.35)
646.06 (2.289
WALKER
TD-CFM
15.08
0.920 (0,029
768. (60.48)
TD DD
GPI
657.58 (1.57)
692 29 (1-40)
637.44 (4.71)
(D-CFM(C
13.57 (3.09)
0.414 (0.019)
179.30 (24.52)
QUADRUPED
TD--0p
(5.1113
TDP-CFM
13.70
0.307 (0.009
164.75
TD-CTM
GPI
669.55 (5.53)
672.46 (3.96)
668.61 (S.71t)
TD CPM(C)
GPI
690.52 (4.31)
697.66 (6:2.)
6066.18 (1.20)
FB
(5.94)
FB
GP
35.23 01:00)
57.51 (1.20)
39.04 ().04
TD-00
312.65((1..1)
(14.25)
633.37 (20.2-)
WAT KER
1D":00
588,
TD-CFM
GPI
506. 06 (15.8
524.34 C6 75
557.24 (17.20)
TD-CFMEC)
513.24007770 (17.71)
634.19((19.0)
1D2-CPM
617.08
647.883 (7.9))
28

---

Temporal Difference Flows
Table 16. Per task results for the quantitative multi-policy experiments.
Walker
Quadruped
Task
Method
EMD
NLL
MSE(V)
Task
Method
EMD
NLLU
MSR(V)
TD-nD
24.22 (0.37)
1,585 (0.021)
494.85 (221.30)
TD-DD
20.23 (1.67
1.394 (0.024)
270.84 (168. 15)
TD"-DD
0.893 10.616
(0.005)
103.71 (34.77)
TD'-DD
135.62 (0.10)
0.921 (0.044)
562.83 (170-42)
FLIP
TD-CFM
16.01 (0.33)
12120 (0.0073
451.62 (61.00)
JUMP
TD-CFM
15.25 (0.12)
0.960 (0000)
365.14 15)
TD-CFM(C)
14.77 (0.39)
0.704 (0.0833
74.42 (13.199
TD-CFM(C)
15.24 (0.13)
0.548 (0.009)
129.02 (23 69)
1D'-CFM
14.81 (0.50)
0.546 (0:0122
73.86 (36. ())
TD²-CEM
15:00 (0.05)
0 3,399 (0:001)
139.10 (9 (9.66)
TD-DD
21,28 (0.97)
1,389 (0.000)
53.28 (20.52)
TD-DD
20.06 (1.67)
1.405 (0.01.9)
273.65 (1922 14)
1D'
102.69 (3,60)
0.516 (0.070)
6.35 (0.58)
TD'
135.28 (0:10)
0.909 (0.049)
171. 76 (18 29)
RUN
TD-CFM
14,99 (0.05)
0.815 (0.083)
80 (51:21
TD-CFM
15.01 (0.02)
0.961 (0.031)
189. 56 (63.62)
TD-CFM(C)
13.01 (0.35)
0.260 (0.080)
32.84 (8.26)
TD-CPM(C)
14.92 (0.17)
0.538 (0.017)
64. (6.77)
TDP-CFM
13:20 (0.30().33)
0.180 (0.071)
34.61 (21.58)
TD2-CPM
14.71 (0.12)
0.351 (0.008)
90.48 (10.33)
TD-DD
21.31 (0.05)
1.482 (0.015)
10933 50 (700.54)
TD-DD
20.01 (1.78)
1.401 (0.033)
1131. .99 (863 23)
1DD
103.72 (1.60)
0.903 (0.067)
115.58 (28 18)
TD'-DD
135,81 (0.19)
0.876 (0.05-)
669 65 (148 88)
SPIN
TD-CFM
15.16 (0.54)
1.020 (0.095)
482.78 (24.82)
STAND
TD-CFM
14.91 (0.10)
0.931 (0.03.)
735.43 (274.17)
TD-CFM(C)
14.22 (0.06)
0.605 (0.076)
170.20 (48.23)
TD-CFM(C)
14.75 (0.30)
0.508 (0.017)
336.02 (16.08)
TD3-CFM
14.34 (0.20)
0:119 (0.056)
197.13 (26 989
TD²-CFM
14.49 (0.20)
0. 0,309 (0:015)
$25.59 (79.53)
Tn-DD>
21.34 (0.00)
1.450 (0.029)
504.94 (210.79)
TD-DD
20.55 (1.05)
1.412 (0.035)
314.53 (64.50)
TDD-DD
103.86 (4.22)
(0.000)
250.96 (79:14)
TD--DD
136.45 (9.57)
0.901 (0.050)
256.91 (38 22)
STAND
TD-CFM
14.28 (0.32)
0.829 (0.3075
1971.68 (1225011)
WALK
TD:CFM
15.06 (0.25)
0.949 (00030)
274.17 (08.60)
TD-CIM(C)
13:43 (0.3-)
0.335 (0.067)
265.00 (12.84)
TU-CIM(C)
15.02 (0.00)
0.518 (0.024)
156.72 (1) (7)
TDP-CFM
13.52 (0.61)
0.284 (0.062)
(66.16((17.5)
TDP-CPM
14.76 (0.19)
0. 0.331 (0.019)
160.62 620.6226 (23 36)
TD-DD
19.30 (0.80)
1.282 (0.053)
622.00 04 (186.99)
TDD
103.79 (3.80)
0.171 (0.055)
425.65 (131:20)
WALK
TD-CFM
14.97 (0.14)
0.787 (0.070)
1344.77 (140.35)
TD-CPM(C)
12.30 (0.25)
0.165 (0.042)
354.40 (1 14440)
TD3
12:63 (0.48)
0.078 (0.072)
301.97 (21.5533
Pointmass Maze
Cheetah
Task
Method
EMD 1
NLL
MSEAY)
Task
Method
EMD
NLL
MSEXV) 4
TD-D0
(2.03
(0109)
16:2.91 (2633))
TD-DD
(0.47
1.358 (0.003)
903.12 (207.00)
TD'-nD
(2U09)
0.63
-D--pD
73.44
0.782 (0.065)
308.54 (30.3K)
FAST SLOW
TD CPM
(2.00)
(2.00)
575.00 (16.51)
FL
TD-CFM
13.06
0.964 (OL077))
911.92 (186.15)
TD-CBM(C)
0.061 (3.002)
0.083 (3.05)
98.04 (5.55)
1Da
0.060
(0,099
61.88
TE-CIMIC)
10.96 (0.58)
0.564 (0.050)
328,99 (27.34)
TD-D0
0.151 (0.007)
2104 (1.19
537.8 80 (HT 50
TD2-CFM
10:95 (0.32)
0.443 (1.047)
222.71 (32.99
0.397 (3.010)
0.60 (0U08)
213.98 (S1.5))
TD-DD
18.64
1.422 (0.002)
678.24 (40.09)
LOOP
TD-CPM
D0088
0.782 (2U05))
225.96 (00.30)
TD3-DD
(6.07)
0.753 (0.007)
215.67 (30.77)
TD-CBM(C)
(2.001)
0.366 (1.09))
86.
(20.49
BACKWARD
TD-CFM
12 (0.33
0.966 (0.022)
381.99 (112.00)
TDa CFM
D.00 (3.003)
0.101 (3074)
102.65 (27
m-CrM(())
(0.36
0.582 (0.005)
230.82 (14.13)
TD-00
(2,05)
(X144)
207.45 CH
TD--CPM
(0.15)
0.476 (176600000
255.25 (37.00)
BOTTOM LEFT
REACH
TD-CPM
0.839 (o.000
0008
378.54 (1T FD
89.56 (0..00)
TD-DD
73.06
17.61 (0.59)
0.742 (0.003)
1,489 (0:00-4)
111.78 (19.51)
87.64 (10.75)
TD-EPM(C)
0.08800000
-0.99 (3.094)
73.84 (1:00)
RUN
TD-CFM
(2.22
1,103 (0.000)
194.36 (19.05)
TD
0042 (3.002)
0.266 (0.021)
14.20 (1.47)
TD-CPMCC)
10:75 (1.17)
0.535 (0.028)
34,90 (21-47)
(2000)
2107 (0.00))
55 (4)
TD²-CFM
10.74 (0.00)
0.445 0.445(00002)
24.71 (10.01)
(3.05
(X.00
1451.00
(342
REACH BOTTOM
LEFT LONG
TD-EPM(C)
TD-CPM
(2.00)
0.078
0.457 (1055)
0.88 (0.000)
2212.63
903.00 (639
(501 47)
TDD-DD
TD-DD
(8.55
73.43
0,802 (0.011)
1.475 (0.009)
90.76 (20.95)
(6.89)
(3.005)
(5.09)
S06. 15 (198.05)
BACKWARD
TD-CFM
12:59 (0:11)
1.083 (0:041)
82.43 (5:50)
TD-00
(2.00)
2,106 (0.11)
194, (75.47)
TU-CIM(C)
(1.17
0.632 (0.008)
30.50 (1.34)
-D--an
(3,05))
0618
112:34 (12.17)
TD'-CFM
11.53 (0:15)
0.634 (0.013)
33.52 (3.75)
REACH BOTTON RIGHT
TD-CFM
01066
(0.05)
272. 15 (60.30
16.77 (0.96)
(0.007)
805.61 (158.80)
TO-CBMCC]
(3.005
0.063 (3.05)
103.89 (10.50)
TD--DD
72. 72.44 (7.96)
0.737 (0.020)
348,70 (18.05)
TDP-CPM
0.07 (0,07)
101 (0.0913
40.28 (13.50
WALK
TD-CFFM
11.91 (0.15)
1.005 (0.041))
1899.15 (151.00
TD DD
0.149 (2.00)
433.81 (97.99)
TD-CFMÍC)
10.72 (0.14)
0.551 (0.024)
277.74 (117.43)
REACH
LEFT
TD-rBMeC)
TD-CPM
(2,00)
(3.061)
0061 (3.003
(2.00)
06311 (3.46)
0,022 (0.00)
453.39 (60.5))
158.15 (35.48)
54.48 (1-15)
TD² T02-CFM
TD-DD
(1.10(
10.66 (6.16)
1.427 (0.0771)
0.464 (0-025)
260.01 (150 430
44 (206.20)
TDP-CFM
(0000)
(T.01)
WALK BACKWARD
TD--DD
(7.4))
0.787 (0.000)
447.90 (40.14)
(.004)
1.994
221 28 (20.99)
TD-CFM
12.62 (0.22)
1.056 (0.00)
1489.41 (90.80)
0,350 (3.022)
0.563 (2.12)
60.57 (AU.70
TD-CFM(C)
11. (1.10
0.621 (0/000)
501.06 (12.67)
TOP RIGHT
TD-CPM
(2075 (2000)
250.T (40.55)
TD²
11.41
0.623 (11.211
585.96 (108.
TD-CRMCC)
0.061 (30000
(x.02
99.79 (1
TDP-CFM
0047 00.0019
131 (0 (0.0005
17.01 (1U22)
(1.09
2,088
800.61 (159887)
(sU0)))
0.500 (300
192.41 (30.20)
SQUARE
TD-CPM
0.05
(2.00)
0,569 () (2.10
281 43 (20.24)
TD-CP)(C)
3U077
0.568
210.94 (07
TD--CFM
0.05 (0.001)
0.302(0003)
186.33 (36 91)
29

---

Temporal Difference Flows
Table 17. Per task results for planning with GPI.
Walker
Quadruped
Z-Bistribation DVZ)
Z-Distiltation (2
Domain
Method
Planner
Damain
Method
Planner
Random
Local Perturhation
Train Disrilaatio
Random
Local Pertarbation
Train Ditribation
GP)
14.00 Bo-ll
14,57 10:00
GP)
(4..7
246.46 (1 631
TD-00
GPI
225
203 46 (64,07)
102.00022200
TD.00
GPI
673.35
(6.00
690.11 00.34)
10M
117 100
GP)
11641
64 2M
228 14 1440
3:9.77 3:9.7702.700 012799
TD" .000
GP)
744.92
756.42 02
745,290-11
(P)
301
158-90 C109
GPI
F-8 10
75
FG 7M
TD- (TMM))
85000
507.88
GP)
TD" CTM
GP)
505 TO 995
302 TV E19
504 (4
TD CPM
GPI
TIN (19
a0.300
TM 45 (44)
FRI
3U5 41
FRI
453 60-99
80,00
456.710144
3024
GP
197
07 (SA 240
BUN
GP
496.35 0130
10:0950
GPI
(10-30
3021 94 230668
326 4T (33.00
TD:DMM
GPI
085.00
(1.80)
45x.4 0190
485 0130
YD-DRMIC
GM
201 21 (64.300
ST (73.7)
YD-CAMIC)
GPU
411.00 (277
NS250
FBI
GPI
20424 13.000
K2 E3
FBI
GPI
913300
(1.90
017244
GPI
GPI
420.30 (179)
100 75
10-0355
GP
848.08 (7.50
TD-C3M
GP)
(1.27
1er 01:20
TD-DIMCC
GPI
555.42
561.7 202
947.29 05.43
TD-CTM0C))
GPI
964.00 (1.3
084891
963.59 01.30
TD' -CM
GP)
923.77 Oluie
0G2 10
TO' cMM
GP)
064.30 0120
961.22
084.55 0.20
(P)
18.58 (0v09
$1.109
GP)
(5.037
658.78
415-2366822
TD-D0
GPI
387.76 881
70011 (11 2000
MAT. (0.109
TD:00
GPI
297
20 45 109
404 (197))
MALK
TD
GPM)
194.45 (1.70
H42 128
8922 022000
WALK
7D"-c0
GP)
457 (23.7)
468.76 (1.0
452 (6246427110
6G8
TD- (mMM))
TD TD-cmste)
GP)
107
GPI
366.05.0541 411
868 ON (14 1-40
8GB 46 (14.44
T07-cPM
GPI
516.07 (1-00
$13.77 61
017.43 541
Puintmass Man
Z. D-Ditiilaatinn D(2)
Domuin
Mothed
Plansee
Z. Dittrbbttion DV3)
Randon
Local Partrrationn
Traie Bitnttatiinn
Randem
Laral Dettrrtation
Train Diarthation
12,
PB
24.9222.00 21
6 30
CM
54 174
345 R2 155771
PART OOO
112 (2000
TD-CEM
FR
GPC
G15
GPS
374 (55.0)
31T 39 AM-
589 GHoil
772 (1)
TD' -c-
7hh C794
107-11
856 69 (199 30
71216 (07.77)
79185
722
400 28 NIE
712 QUAGAe
102382
115.00 E117
341.00
31
T0:00
1000
10/00
GP
875.00
PELP
TD-30
CALDO
10-3M
CP
GPS
932
792.00 DOLIT
RATOMARD
546 12 (15 di
S (64
AND J5 T5L10
GP
7H0 2e ing
700 19019777
684.90
0OL 25 O1
12.90
TEACH
TD-30
T0-CMM
CM
75
20-C0M
243 SOAUTT
2206 05
397-C2M
600 05 10
age S8
TD--crM
412.1 01.70
407 26 COLIO
454-90 506-30
GPC
REACE BOTTOM
LEST LONG
10:05M
TD-30
390
S1L (1
840
10
430 AT RUR
MEN
12-30
288222 (100
1.11
546.221
GPC
ESL 82 0L.
000 AR0
000-44 FL00
BASKWARD
20-C0M
250 14 140000
5O0 30
92 45 (15.00
3 15 C01
2001 BEYGORG
FB
FR
GM
0
335 46
10420
UFF
BOSTON ROURT
TEP.0
GFC
396.78
19.55 BOAAT
12:30
99% 4544
CPC
ENC
VALK
m0/00
1VI 16 Vavae
665.00 L2V
CPC
RAN DEVI THS
1864
375
TE-CTMOC)
107-CTM
11 62
BASKWAND
VALK
30-00M
T=/000
12:30
450.14 (191-10
496 20--2-0---000 F13
546.00
350 JA E--UM
982 VILLE0
755 (1)
TD-CTNICI
CMM
164.711150
963399 01 20
9822 19 (1LE3))
155
082.51
1802 01110
REACH
T27.00
T0:CFM
823.30
ROLST W8
NOW 33 C3.06
TROJINGL ED
Fil
A3R 96
GM
164 50
T0:00
MR. TR
100 A
12P40
695.887750
TD-CFM
UFF
AI.2 C2.70
CPC
93036 12 1S
30

---

Temporal Difference Flows
7 = 0.8
7 = 0.9
7 = =0.95
7 = 0.98
7 = 0.99
TD-CFM
TD-CFM(C)
TD CCMM
TD-DD
ID2-DD
VAE
GAN
GROUND TRUTH
Figure 5. Qualitative samples generated with TD-CFM, TD-DD, VAE, and GAN methods for various discount factors y on the LooP task in the POINTMASS MAZE domain. The last row depicts ground truth discounted occupancies.
31

---

Corollary Let {mi" (") be the sequence of probability paths produced by TD-CFM, TD-CFM(C), or TD2-CFM starting Temporal Difference Flows
E. Theoretical Results
E.1. Proofs of Main Results
Lemma 1. Let Pr. be a probability path for P generated by vector field in and p(") be a probability path for p"m"
generated by 0,() such that pu Po my. For any t [0, and (s,a) let of 1) s, a, be the solution of
(())
arg min (1 B) X.mpe(+1a,a) [|]((X)) ox s, a)]l
e: 2: -Rd
t, R. ) (-b,aa) [|[({t)= 06(") (x, aa)l
Then ui (n+t)) induces probability path m 1n-1) such that mg (n+1)
ma and m (a+1) m= (m)
Proof. By Lemma 4, we have that
1 (xr 18,0) (1= a) s,a) + "pt (z) (xlss a) (x s,a)
mi m( (x)(s,))
(n+ 1) (xrls, a) = (1 - ( + YPt (x|s, a). Lemma 3 implies that m
(n+1) is the probability path generated
where my
by et (n+1) is casy to that my
(a+1) mf+t1; mo since po Po
(nn) mo- Moreover, since pi P and p (m)
pmm (n) by
assumption, mi (n+1)_ = (1 - P +
p (m) = T m which proves the result.
Theorem 1. For any n 2 1, the probability paths generated by TD-CFM, TD-CFM(C), or TD²-CFM satisfy
mi
(n+1)
(x ss
(B7 (1)) (I s,a), At [0,1]
where m, P and P,(x)s, == Paj (2) x P() Is, d)dr For any t E [0, 1), the operator B is "y-ootrraction in I-Wasserstein distance, that is, for any couple of probability paths pt, 96-
sup s.a
W1 ((B p.) (- s, a), (B74) q:) (- s.a))
7sup Wn (pe(: s,a), al s,a)).
s,a
Proof. To prove that the iterates of the three algorithms satisfy Bellman- like update through the operator B7 only need to apply Proposition 3 for CFMM Theorem 5 for TD-CFM, and Theorem 6 for TD-CFM(C). That B, is 7-contraction in 1- Wasserrstei distance can be seen by applying Theorem 4 with k 1.
from an arbitrary vector field (0) Then,
lim mi (n) mT BiTTi n-+00
where m, the unique fixed point of B,, and m. mnc where mpC(- s.a) S 21) m= s, is the probability path of the Monte-Carlo approach (MC-CFM; 6).
Proof. That B7 has unique fixed point m, to which every sequence m. converges to is a consequence of the Banach fixed point theorem applied on the space of all probability paths m S x A + (RR)) equipped with the sup-1- Wasserstein metric. By inspecting the definition of By, it is easy to see that my (I pp P.. Since (r(ass, Pqp ((rl)) a)dz1,
m a) [(1 (pp))-PP](x(s,,. a)
- dæ = MC (s, a).
(4) (م,ع)
32

---

Temporal Difference Flows
Theorem 2. For any n 2 and E [0, 1), assume that my (m) (x 8, ppp am)m ( 8, a)dd,, then
T CFM
2² "To
(X)(x())))
Proof. See Theorem 7.
Theorem 3. For any n 2 1 and & & [0, 1l, assume that (r) (x 8, 8,a) - P | Io, (o, 21) (o 21 | 8, addrddz s, then we obtain
TD-CEMMC) = T CFM
+
FE [Tr(CON (X, A) 21 (x]z(()))]]),
where Xi). Furthermore, if we use straight conditional paths, i.e., Xx tX + (1 1)Xo, and the linear interpolant X, does not intersect for any s, o. s', then o (D-CIMIC) -0
TD2-CmM
Proof. See Theorem 8.
E.2. General Results
Lemma 3. Let ot and os be vector fields that generate the probability paths pt and rÊ. respectively. Then, for any y E [0, 1].
the mixture probability path pt (1 -))p + YP? is generated by the vector field
0t (1 7)pp o 7ppr0
(27)
(1 ?)pt YPI
Proof. Since of (resp. u generates pi (resp. př), know from the continuity equation that:
op: div(p!u)).
dp
div div(p))))
Ōt
at
where div denotes the divergence operator. Then, by linearity div,
Öp: ((1 r)pt + xpf)
Ot
(1 d) 0z)) ((di(pp
div
((1 - 7)p + v3)
(1 - 7)p uppr? ((1 711-))p + tnrt)
div
(1 7)pb + ypr
(1 - 7)pp vt + qp, o?
div
(1 7)pb Tp? Pt
div
Hence, Pt) satisfies the continuity equation, which implies that generates pt-
Lemma 4. Let ot and vi be vector fields that generate the probability paths pt and př. respectively. For E [0, 11. the
vector field v (1-r)phut =ppi= satisfies
= arg min {11 7) Ez, <ph [1(((x))) - 1/ [+))) + ExswEe [llve(xe) - -2 (x(r))
Proof. Let (v)) 7) Baawn! this quantity evaluated at some point x:
vz (2+) + Ba (1o 12 (2x))1))) The functional derivative of
v.l(v)(x) - 7)r (.. (+)) oz))
Setting this to zero and solving for (xx) yields the result.
33

---

Temporal Difference Flows
E.3. Analysis of TD²-CFM
We study the learning dynamics of idealized variant of TD²-CFM which minimizes the flow-matching loss exactly, Starting from an arbitrary vector field iy (0) at each iteration n 2 0 we compute
(++ (1,,a) E arg min ((n) "D2 CFM (1, (,,a), a), w.d
(28)
where
(())
CD- (t,,, = (1 (1- 8, +
(tt, 8, a) = E (1 ( (x ( [S)
s,,) == E E (o)) (1 (1.))))))) [|o(Xlls,,) (X, ( ((S))\]
and my (xlss,) (2)s, a) is the probability path generated by "((" (x)ss
Lemma 5. For any 2 O, the vector field minimizing (28)
(1+1) (x s,
1--7) f Tune x1) x1) (2) + 7BBss (x|"s, (z')", (m) (x|S", +(S))
m, (a]ss
where we define mi x1) P1)P(x1le (s, a) Moreover ui (n+1) generates m, (x+1)
(x+1) (z)(s,a) a) = (1 - (allss)) 7 1
(1) In (x|S", and Pr(x]ss, a) = fpup(4
Proof. By Theorem 2 of (Lipman et al., 2023), we have for the first term in 2
Volll,s,,) Volxx [1 (X,1s,,
where P((x|s,a) a) :=
(x x)) a)dz (x)ss = Pls (s,ader Similarly, we have for the Pe 8,0)
second term:
V Volll,,,a) = Vol. x x xm (1m0) (1 (x) B ((xelss)))1
where (m)) p= (n) and B
(mm(")) (00)
P p'm(")
Therefore, TD-CEMM (1, a) is equivalent, in term of gradient, to mixture of two marginal flow-matching losses, which implies that o, (n+1) has the stated expression by Lemma 4. The fact that it generates m. (n+1) is consequence of Lemma 3.
We then define the following operator to characterize the iterates of TD2-CFM Definition (Bellman operator for probability paths). For any t E [0, 1], we define the operator B7m := (1-7) 7 m where Pa(x]s, a) : Puju(C
P (x w1) P(x [s, a)dx
The following observation is then immediate from Lemma 5.
Proposition 3. For any n 2 the probability path generated by TD2-CFM satisfies mi m (n+1) (x)s,) a) (") s,
(ntt1) (x(s, (B7 (")m(())( (x s.
where BT is the operator of Definition 1. Moreover,
Theorem 4. For any E [0, 11, the operator BT of Definition I is 7 1h conttacctinn in Wasserstein k-distance, i.e., for any couple of probability paths Pt. 9t and lo E (0).
sup Wi (BTT)) 8,,)), (B) ()) 8, a)) & 1/k sup We (pe(- a,))).(( 81,
4.4
8.0
34

---

Temporal Difference Flows
Proof. Recall that the Wasserstein k-distance between Pt and qt induced by metric d is defined as
inf
where C(pa(-ls, a), (-(-ss,)))) a)) is the set of all couplings between the two measures. Now take any coupling ('(1ss,)) C(pi(yls, qo-1s, a)) for any Then, the following quantity
(z, = (1 - P(x)s) - v) + (2
a valid coupling between (BTT.) (- s, a) and (:) a. fact,
/ 0(1, (1 -)
7) / P(rl)s,))((r - y)dr + ? (1 7) P(yls, +7 Ey B F(c.vls' dx
(PTT) y s, a,
= (1 (1->))pP))) P(yls, a) + +7 + / N(a.vls' "ec.)uss==ss')0000 =s
= (1 P(yls, + E. B (aluul''==<
(T*x) (vls,a))
Analogously, we can prove that / (z, uls, aydy (15 (c) a). Then,
Wi ((B7 ( a), (B) s,)))
inf s ((t
E -() [d(X, 8)q Y)
(11 ) (x d(x, V) + 7 (x, ((, 1/2
=/1/B p1
Since this holds for any coupling i'(-ss,a) right-hand side, so that
(.-ls,, a)), we can take the infimum over all such couplings the
[recton inf Foxxyila(x r)91) 1/k
Wi. (), (.),a)) S (EemP((le.))
= =/ (EsnP(()),))
(1rlsssss' (())) ((())))))))
S 2/k sup WA (pr(
8. a), (:(: (),a))
8.0
Taking the supremum over (s, a) of the left-hand side concludes the proof.
E.4. Analysis of TD-CFM
We study the learning dynamics of an idealized variant of TD-CFM which minimizes the flow-matching loss exactly. Starting from an arbitrary vector field uz (0) each iteration 2 O we compute
(m+1) (-|s,a)
arg arg min An) st, (.)
s, == x1 (T=m(()) (s.)) (s.m))")ss)))) (x [lx((x) [llx(X))- (x,(xx)) (29)
where m (x) (z|s, a) is the probability path generated by (n) (x)ss,))
Lemma 6. For any 2 0, the vector field minimizing (29) is
(2) ss,a) (n+1) (x s, f - (2/1))) (xp x11 (T) (r))
dx,,
m (n+1) (n+tl(
where mi (n+ (n+1))
(2|8, a) == prpu (x | (m=)m 8, s,a)d Moreover of (n+1)
generates m. (n+1)
35

---

Temporal Difference Flows
Proof. Note that (29) a standard flow matching loss for the target distribution T Tmm"" (m) The expression of t/ given in the statement exactly the vector field obtained by marginalization of the conditional vector field uenn which we know to be the minimizer of the loss from Theorem 2 of (Lipman et al., 2023). The fact that 2 1) generates my consequence of Theorem of (Lipman et al., 2023).
(n+1) !)
(ar (s a)
(a+1)
is
Lemma 7. For any n 2 0, the probability path generated by (29) satisfies mi (1=1)))) (x]s, = (T*m ( a).
Proof. This is immediate from the definition of conditional probability path, as we set Pip | x1) = 8(x - x1) by construction, where 6-) the Dirac's delta function.
Theorem 5. For any n 2 I. the probability path generated by (29) satisfies
m a) = (5, (3) (x)ss.)) (al)s))
where B is the operator of Definition 1. Moreover, if the initial vector field o satisfies
of" (x s.a) = f (2)pa))))) Petu (Ir
x ) (0) (01 8, a) dæ1.
(3x) ( a) m
(0)
with m (0) being its generated proability path, then this result is valid at all 2 0.
Proof. We know that, for all 2 0. uitt1 generates m; ((++t) (Lemma 6) and that mi (n+1) 7m
(n) (Lemma 7). Note that
mi (n+1) is written as a function of mi only, i.c., at each iteration we keep only the distribution generated at time f = 1 (n+1) function of m. (n)
(m/") and discard the associated probability path (ml (m) for t < 1). We can however express m;
thanks to the linearity of the Bellman operator and the definition of marginal paths. For any n 2 1,
(xl | a,a)= (n+1)(xx)
= f pipl an (T)(TT ((("))( (xr a,))d
me
= -fmuler panll and (1 (1-7)) P(r1 s,)a) + P( mi (m) (or I (6,666)) dxi
= (1 (1->) f Pin(al P1) s,a)dr + r BB sa)
/ punle x1)m( (m) (x1 | s )dzs
= (1 - 7) f /p/1(00 P(n s, add + 7 Bl a0)
/ Paulax (T= (2) ss, x (s'))
= (1 ->) / Ppi 21) s, s,add + y BBsppp s0)
mn (x |s, ((3)))
= (1 - P(xls))))))
(B7 (n) (z s.a). + Bp m (n) (x s xs'))
This proves the first part of the statement. For the second part, we only need to prove that the result also holds at n 0. Note that the assumption on of" implies that m, (x (x(() s,a) == pin(" 21)m (9) (x1 ss s,a)dz Thus,
(())
mi ( (I s,a) / pud (x (I x1) (T (n) (2) s,a)dr
= /maleel xi) (1 (1-7) s,a) + + Ty +7
mi (xr 18', (((6'))) da
= =11 / ||1(=|21))(x) P() s,a)dz + 7 Eg B(
/ Pati(ze mz)mi 10) (21 ss', "'')(s'))dz)
(1 -1 / (m](=|)))(( s,a)dz + 7 (()
my 0 (x) ( s(,, (B, m3 ( (z s))).
36

---

Temporal Difference Flows
E.5. Analysis of TD-CFM(C)
The idealized update of TD-CFM(C) is, for any n 2
Di ((+t)) (11,a) a) E arg din)
ID-CMMcc s, a). where
(30)
((a) ces (1, 8, E (Xn. (..) (1) [Xxx uojo,1 (X+ Xo, Xo, x
and [(a)( (- a) is the coupling between mo and 7 m=
(m) while P, 11011 are such that (z) 20, 20,01) generates
P (x 20, T1), (1 20, = 8z, (z), and
Peji (xr x11 = / Prjo. (Ir Io )
(31)
Lemma 8. The coupling (()) 0,1 (- a) satisfies
(m) (2o, x1 s, = (1 P(x ss. (.))) + P()a))) mo.1 (ro, xı |5 ((S))
(10) where m0,1 (20, 11 8,a) = mo(on)) of (20l,,)) (x1) is the joint distribution of (Xo, X1), ie the endpoints of the ODE.
Proof. For any To. I1, we can write r (zo. x1 s.a) = (1 | s, a, mo)oo())), where r!")) is the corresponding conditional distribution. By definition, we have
r (I1 3, zo) = - P(x s, + + +BBor (ooo
where w(nh) is the flow that generates mi"). Multiplying both sides by mo( and using that m (m) ma(xo) 66 w/ ( (zols.)) (21) concludes the proof.
(Io (,
Lemma 9, For any n 2 0, the vector field minimizing (30)
1x (x s; / u1 za, x0, (x
zo, x1
( To (ro, x1 18, a)
drgdæ1,
m,
1)
(x a)
(n+1) (z s, s, = J P (x zo, x1 ( (ro, dddzodz1. Moreover it (a+1)
(n+1)
where m
generates m,
Proof. Note that (30) is a standard conditional flow matching loss since 1 Fy. r1) generates Prjo. (x (o. and Pijo,( (z Fo. 21) = 8z, (1r). The expression of 0; ( (x s,a) given in the statement exactly the vector field obtained by marginalization of the conditional vector field urjo1, which we know to be the minimizer of the loss from Theorem 2 of
(Lipman al., 2023). The fact that et (x+1) generates m,
(n+1), is consequence of Theorem of (Lipman et al., 2023). O
Lemma 10. For any n 2 the probability path generated by (29) satisfies m (n+1)(x s, (")m (x: s,)..
Proof. By Lemma 9 and the fact that Pijo, (x: To, 8z, (z).
(n+1) (x s,a) == / p1jo, æo. x2 ( (2v. s, adddddr
riul ( s, s,a)dro
= (m"
37

---

Temporal Difference Flows
Theorem 6. For any 1, the probability path generated by (29) satisfies
(n+1)
x s, m, ( (x | s,
m4
where Bi is the operator of Definition 1. Moreover, if the initial vector field ut satisfies
1 4 (x (,)
/
uejo.1 u1 .1) Prjo.1 (x Iu, T1 jm (0)
(To. s,a)) drgde1,
m (1: 8,0)
with m (0)
being its generated probability path, then this result valid at all n 2
(n+1) Proof. We know that, for all 0, U. u+ generates m, m(=+1) (Lemma 9) and that mi [ (n+1) is written as a function of 0,1 only, we can rewrite as a function of my (n) thanks to the linearity of the Bellman operator and the definition of marginal paths. For any n 2 1, By Lemma 8,
(n+1) mm (Lemma 10). While
my
(n+1)( (x 8.a)= == / (1 zo, 2ı J (so, 01 addxodr
my
(n) I s.a)ma(zo) + (gr mo. (zo, x1 15', 15',a()')) dæødar = -f Pajo, zo, 21) (1 (1-
(1 -)) / P1 zo, 21) (x
+ Bu [ //aao.10 æo, 21 a1) (m) (zo, 21 SS (SS))ddz
(82)
By (31),
(1) = peul(r ( P(ri a,,)ddi = Pr(ar s.a).
For (ii), by Lemma 9, we have ml() (x a) [ (= 21) (n- 1) (20, 21 s, a,))drd Vn 2 0. which implies
(n)
(n) mo,1 (co, 2n 1s',m(8)) = r
(20, 21 "d0
(8)))
(n) (Lemma 9),
Therefore, again by definition of m;
[ |prootoe
( = E P()eaa)))))
Proo1(1)) (x wa, xn m (ro, æ1 18,
(m) (x ss ((('ss))
Plugging the expressions of (i) and (ii) into the one of m m. (m+1) (xc a) yields the first part of the statement.
For the second part, we only need to prove that the result also holds at n Note that the assumption on E (0) that mi (0) (x s, = f i (x zo, )m (Zo, (2.. 21 s', (ssd)) Thus, using the same decomposition above, we have
(0) implies
[ P xo, x. m
m (1 (z s, Pz ss)a) + B.
/ (zo. x1 S',T(s') (sss))dzodz)
= =1177) s, + E. Bamp((la
mi (x s (1))
which proves the result.
38

---

Temporal Difference Flows
E.6. Variance Analysis
Theorem 7. Let us define the random variables
O,0 Xx, x(_)) = (1-2/Voo((X(sss0))
- (Xe((()))))
+ ( (s a; (vx()) - (x, (x;
s,0,, X1, (1 (V (1, (1x(X))) (x,(''))))
(X,s( (M) (u,(X, tinn (X,(x))))
where t u(0 1)). (s, p,s' P()ls,a), Xt (-)(sss"),, x mi m s' xss))), ml | (s'))
and X, ~ Pan (1) Then, 9yo2-cra and 9TM-CPM are respectively unbiased estimates of the gradients Vol crx (0) and VOlTD-CMa(.).
Moreover, we consider their respective total variations defined as:
2 TDR-CFM Trace (oov 1, ( X2, x(1)))
TD--MM = Trace Corr Lis,a,o XenX (tD,cau(t,,,,,58 X1, x+)))
and we assume that (n) (x s. a) =
Pin( = mm)mm (r1 8, a)dx,, then we obtain
2 TD-CFM = = CFM + Bi,s,, [Trace (Covx, [Vovi(Xels,ar a; 0) (T ueui(Xe & 1 xx)])]
Proof. Recall the TD2-CFM and TD-CFM objectives:
(rm = (1 -1--1 B) "x.saa"""Xexpniii Xempin((s)) [|oc(x,ls, a:0) - (x,s'')))) [2]
y t,s,a,'' "as,a,'i,,mm Xm (h (1e (f (1o,s a: - (X,ls' (ss))|]
[| a; - un (x,1''))) 1/9
(TD) = - E Xempen
71E t., XX Xi (est
( -xn
Computing the gradients of these quantities w.r.t is easy to check that 9ID2-CrM and frD are their unbiased estimates.
Let us now analyze the total variation of these estimators. By assumption, we have m, (n) (z s, Pai =1)m{" (n)) (ar I
a)dæ1, which implies that (n) and X, follow the same law. Moreover, we obtain the following identities:
n" (x ((')) Exyla [ugn (x: X1)).
InD (1,8,0, s. Xt, X,) = Ex. (or (t, X1, X.)),
Ex (x (169)()))) on s, X Xi) E Ex (16') 9ro-crM (t, s, a' X X1, Xx) Xe<Pen((X1))
)mm m! m (a)s.a)
where X1 x, 5.
the posterior distribution of X, given x and 8'.
To simplify notation, we denote by Y the random variable (1, 8, a, 8', Xz). Using the decomposition of variance into
39

---

Temporal Difference Flows
conditional variance, Var(X)= E[Var(X][)]) + Var(E[X]Y]), conclude that
OTD-CFM Trace (Covy,x,,x, LOTD-CEM (Y, Xx)l)
Ey,x1.x, | X1, X2) X WNN (Y, X1, (1, Xx)] [|]]
Ey, [| lerv-ces (Y X1, X,)] Y Lor (Y, Y X1, Xx]] |9]
Br.x, [ExIXXX, || OTO.CRM (Y, X1, X1) ExIYx, WTD- (Y, X1, X,)] [']]
Ey,x [| (X, Ey 1010" .om (Y, x,)] 1/]
7 1y,, Ex,,,x, |/l Voo
(xn X1) Ex, [Vou((Xels, a) wan xxa) [[]
CFM + Ey,x, Trace (Covx, Vou (X, xx)])]
Tpt-cm + 7 B,,,x, Trace (Covxilaa,x) [oore(X 0) uen xi)])]
Theorem 8. Let us define the random variable
(1o.CrMic (t, Xo. X1, X1, := (1- (v,(x,,s, ur (X,lo s'))
+ +~Vovi(X,ls,,a (v (vr((x,s, - "non No, X1))
where 1 ~ u([O, 1)), (8) ~ p ~ a), ~ (ss() (No, ~ mor(( 8 (s')) and X2
P X1,). Then 9rv-crste) is an unbiased estimate of the gradient VolTo-rrmc)())) Moreover. if we consider its total variation defined as:
UTD-CIM(C) Trace (Cov t, ", Nn, (tt, Xe, Xo X1, X1, Xi)
and we assume that m, (m)) (x a, = ʃ 1 zo. 20. x1)m (m) (ro, 21
s, a)d then we obtain
GTD-CrMeC) TID--CMM + Be, [Traee (oo (Xe,s a; 0) "eeo,, Xo, Xi)]]]
Furthermore, we use straight conditional paths, i.e., Peju, (xr x1) = 8(1x) (1 1)xo x)), then
OTD-CFM(C) S0Tp² CDMM-CFM
+ + sup (,)((x)s,,;0)) E [[|X] Xo Xx [X, - [3] 1,a,a,
In particular, when the paths of the linear interpolation X₁ do not intersect for any s,a,8' we have X X1. [[[X - Xn (xxx x, [X1 - -Xo] [2] = and Tnn-CeMc)) = O
Proof. The first two statements can be checked by repeating the proof of Theorem 7 with conditional paths PHo,1 and vector fields 1. Let us thus prove the second part. We know that the flow d (2 21) that generates the the conditonal path Pejo,1 (2) (o, 21) 61m.+(1- (x) is (e 21) = tr (1 t) Its associated vector field uejo, 1 is thus
d
u (or 21 220, x1)
(x(0 200
40

---

Temporal Difference Flows
Theorefore, denoting Y (1, 8, a), bound the second term in the decomposition of TDD-CEM(C)
By,x, [Trac (Cov(Xo.X))IX.., (Xrs(X uz/1 (Xt Xo.
Exo.X (x a; 0) (o, Xx) ExxxX,/XX, [Vou a; a) (X, Xo: x1)] |=]
-,
Xo, Xi) [1]
S Er,x,
(x,8(X, a; Xx.
|i Xo, X1) X, [
[xo X1 -Ex.,X, [X1 Xxo [19]
[X,,(X, a; 0)1 2 Xo.XX
< sup [ a; Brx [ExoxxiIX.x, ||x - 10--1--1 - -xx.xxxxX (X1 - Xol |"] 1,,,,,
This proves the third statement. For the last point, simply note that the paths generating X, do not cross, then the distribution of No: X1 X, is supported over a single couple (Xo, X2), which means that its variance is zero.

---

Temporal Difference Flows
E.7. Transport Cost Analysis
Theorem 9. Assume that mi (m) (x 8, Pup (1)m (3) (x1 I 8, a)dx,, where Pei(- | = -N(11,, (1 112 is
Gaussian path. Then, the conditional paths built by TD-CFM(C) and TD²-CFM to generate mi (a+1)) = TTmmi (n) smaller transport cost than those built by TD-CFM. Formally, for every 1, 8,0,
induce a
Et,,,a,s' , Xi00mm, (0t) (xol)aat)))) - -xl1 - L, Xo Xi ( [ Xoll3].
Proof. The paths generated by TD-CFM(C) and TD²-CFM induce the same transport cost since both algorithms connect the endpoints of the ODE path m, (x) in the bootstrapped term. Hence,
, No (101)11-))))))) 17 ) (xp1)/00)))) [[|X] Xool²]
(1 7) Xo |||s - xoll + y a Xo [1 [ (Xo ( s(' - xool]]
(a)
(1 7) No [|ls' -xol + yE T
[II
(m) (w) (Xo s' x( ||]
(0)
E,, [||s Xool?] ]|] + E (10 (Xo 18 x(, di
@ (1 7) E No [||s Xoll +y / (ieaa"
01 (h) s x (d)))| dt
(4) (1 7) B, No |||s Xol1 +7 f B X (-Io(_m())
Ex dtan (X+|xii
dt
(c)
[|5' Xxoll] + / X
E Exiloxx (un((x > |1]]
S (1 B)
@ (1 B), [||s - -xoll"1 +Y
x(xx(((()))))x(( [||411(x,/10 dt
( (1 E, No -xoll]] + +7 f (
|unj (tx, + +1 -)) xx) 1 dt
( (1 7) El [||s - [Xol11 +
/
||x0 - xoll
Xi - -xooo|l xoo1
7) Er,,,,,''' [[]s' - [ + yE L. (1e
E
[||X Xo.1
Laaas No X1
where (a) uses the definition of flow as integration of a vector field, (b) uses Cauchy- inequality, (c) uses that
(n) my vi(m) is the pushforward measure generating mn(") (n) (d) defines X1 x, (a jmm("))) (X1ls' as the posterior
mi (x/s,a)
(x) is in marginal form by assumption, (e) uses Jensen's inequality, (f) uses the distribution of X, given x, s' and uses that vi
Tower property of expectations, (g) uses the definition of Pıja and the corresponding innear-nterpoolation flow, (h) uses the definition of unj, (i) is trivial, and (j) simply combines the two terms using the definition of Bellman operator T*
'Recall that, given a marginal probability path m (m) 8,a), the conditional probability path built by TD-CFM(C) and TD2-CFM generate 7 m" is a linear interpolation between noise Xo ~ mo and X1 ~ (1 (1-7)6,- + ryf"") (Xo s (ss")), while the one built by TD-CFM is a linear interpolation between noise Xo ~ mo and sample X₁ ~ [T [. 8, a) from the target distribution.
42


Reviews:
[{"review_id": "v8KXn1yiBW", "reviewer": "Reviewer_uS9Q", "summary": "This paper introduces a new family of algorithms–TD-CFM, Coupled TD-CFM, and TD^2-CFM which propose a Flow-matching-based generative modeling framework to learn to sample from the discounted successor measure. Learning the successor measure rather than predicting long term rollouts prevents the usual accumulations of prediction errors. The paper proposes several extensions of TD-CFM to coupled flows, analogous denoising diffusion models along with an interesting variance reduction technique using boot-strapped updates (TD^2-CFM). Experiments on several simulated control tasks show that these methods yield more accurate predictions and better value estimates improving planning performance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> **Comparison to direct Q-function learning and finite-horizon planning methods**\n\nWe 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:\n\n\n* **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.\n* **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.\n* **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.\n\nTo 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.\n\n|  | FB | TD²-CFM-GPI | MPPI |\n|---|:---:|:---:|:---:|\n| Cheetah | 479.35 (14.56) | **693.63 (5.50)** | 541.22 (5.28) |\n| Pointmass | 472.45 (14.40) | **800.99 (8.56)** | 286.43 (54.95) |\n| Quadruped | 627.28 (1.98) | **695.73 (2.07)** | 156.80 (122.89) |\n| Walker | 526.66 (5.94) | 627.63 (7.97) | **658.15 (21.46)** |\n\n> **Run-time analysis and computational cost**\n\n\nThis 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.\n\nTo 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.\n\n|  | NLL↓ | EMD↓ | VF Error↓ |\n|---|:---:|:---:|:---:|\n| 2 Steps | -0.48 (0.21) | 0.076 (0.003) | 379.52 (81.75) |\n| 5 Steps | -2.75 (0.15) | 0.036 (0.000) | 23.82 (2.05) |\n| 10 Steps | -2.85 (0.17) | 0.025 (0.001) | 7.71 (2.75) |\n| 20 Steps | -2.99 (0.04) | 0.0218 (0.001) | 4.40 (0.82) |\n\n---\n\n\nWe 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.\n\n### References\n\n\n[1] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency Models. International Conference on Machine Learning (ICML), 2023.\n\n\n[2] Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One Step Diffusion via Shortcut Models. International Conference on Learning Representations (ICLR), 2025.\n\n\n[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.", "claims_and_evidence": "This paper claims that \n-The proposed TD-flow framework can learn a successor measure that more effectively models long-horizon state distributions.\n-TD^2-CFM can reduce gradient variance compared to standard TD-CFM by using direct velocity matching under certain conditions.\n-TD-flow translates into better downstream value estimation and planning performance in simulated tasks.\nTheorem 1 shows that when the TD-CFM is solved exactly, it forms a contraction (in the 1-Wasserstein distance) and Theorem 2-3 shows that variance is reduced through a tighter coupling of bootstrapped targets. Experimental results support improved value function estimation and planning and generally support the claim the TD^2 can improve performance via variance reduction.", "methods_and_evaluation_criteria": "I think the proposed method is a very interesting solution to efficiently learning successor measures. I'm not exactly sure if successor measures are a necessity for the task of planning compared to other offline methods using direct Q learning or simply finite-horizon MPC. However, their benchmarks directly quantify Q function estimation accuracy and improved planning with generalized policy improvement over finite families of pre-trained policy and show significant improvements.", "theoretical_claims": "I briefly reviewed the proofs, but did not verify them carefully. Many results follow directly from CFM. The assumption of exactly solving CFM at each iteration for Theorem 1 is reasonable since we are training in the offline setting and CFM objective can be efficiently optimized if the dataset is sufficiently large enough.", "experimental_designs_or_analyses": "The experiments are quite thorough. I would like to see some run-time analysis for training the proposed algorithms and baselines. It may also be interesting to better how this planning with the successor measure compares to other finite horizon shooting methods /  model-based planning methods if applicable.\n\nIt could be useful to see an ablation as dimension increases or in the presence of noisy data.", "relation_to_broader_scientific_literature": "The work builds on several recent advances in generative modeling for reinforcement learning, including generative horizon models. I believe it is a particularly original and interesting application of flow-based generative modeling.", "essential_references_not_discussed": "I'm not familiar enough with offline modeling of successor measure literature to suggest additional references.", "other_strengths_and_weaknesses": "Strengths:\n\n-The paper is well-written and I found the ideas generally interesting. It is my impression that the integration of TD learning with flow matching to directly model the successor measure is quite novel with several improvements and variations already proposed here.\n\n-Theoretical analysis provides insights into contraction properties and variance reduction.\n\n-Comprehensive experimental evaluation over several downstream application of successor measures.\n\nWeaknesses:\n\n-Even after the extensive evaluations, I don’t have a good feeling for how this framework compares to more direct methods for learning Q function from offline data or model-based finite-horizon planners (e.g. shooting methods, cross entropy method). Even if there isn't a directly comparable setting, its worth discussing their context more and why the added complexity of learning to sample successor measures is useful.", "other_comments_or_suggestions": "A typo on line 190, left column in the equation denominator.", "questions_for_authors": "1. How does this framework compare to more direct methods for learning Q function from offline data or model-based finite-horizon planners (e.g. shooting methods, cross entropy method) ? It could be worth discussing the context of these other methods more and why the added complexity of learning to sample successor measures is useful in practice (besides controlling long term errors better).\n\n2. How does the run-time in training and inference compare between the proposed methods and baselines? It would also be nice to get an idea for what dimension the variance and iterative training scheme becomes an issue through a simple ablation study.\n\nAfter understanding these points, I'm open to increasing my score.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7GngUIKUol", "reviewer": "Reviewer_HuPg", "summary": "This paper introduces TD-flow, a novel approach to using the Bellman equation on probability path with flow(score)-matching techniques for generative models. The approach outperforms the existing ones in terms of more accurate generation over extended horizons. TD-flow is validated across diverse experiments and domains. The paper also contributes theoretical analysis.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> **What benefits does TD bring to flow/score matching?**\n\n\nThis is a great question, and we appreciate you highlighting it. We’ll be sure to clarify it in the revised manuscript.\n\nTemporal 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:\n\n\n* **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.\n* **“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.\n* **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.\n\n> **What do Theorems 2 and 3 imply intuitively?**\n\n\nThe 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. \nIntuitively, 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.\n\n> **Why does Eq(1) hold?**\n\n\nBy definition we have,\n$$\n\\begin{align*}\nQ(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] \\\\\\\\\n&= \\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 \\\\\\\\\n&= \\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)} \\\\\\\\\n&= (1-\\gamma)^{-1} \\mathbb{E}\\_{X \\sim m^\\pi(\\cdot\\mid s, a)} \\left[ r(X) \\right]\n\\end{align*} \n$$\n\n\n\n> **The MDP definition is different from the RL community, where γ acts on reward and value functions.**\n\n\nThank 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.\n\nIn 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.\n\nSince this interpretation doesn’t appear frequently in standard RL formulations, we’ll expand on it in the revised version to help clarify the connection.\n\n---\n\nWe hope this response strengthens your overall impression of our work.", "claims_and_evidence": "Yes, the claims are supported by theoretical analysis and experiment validation.", "methods_and_evaluation_criteria": "Yes, the evaluation is reasonable.", "theoretical_claims": "I roughly went through Section 4 theoretical analysis.", "experimental_designs_or_analyses": "The experimental design makes sense to me. And the results look promising", "relation_to_broader_scientific_literature": "This paper is very relevant to the generative modeling domain, which is also obvious by its name TD-flow, using TD learning (concept from RL) to aid the score/flowing matching in generative models. I feel this paper can/will receive a lot of attention from the domain.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n1. I would say this paper is well-written. For someone like me with limited knowledge of generative modeling, it is easy to follow and effectively presents the key ideas.\n2. It studies a very important yet hard problem, accurate generative modeling for long horizons. And the experiment results look quite promising. The RL community needs an accurate world model for sure, which could facilitate the development of model-based RL by resolving the modeling errors and biases.  \n\nWeaknesses:\n1. It is not clear to me what benefits TD could bring to the flow/score matching and how these benefits help with accurate generative modeling for extended horizons.", "other_comments_or_suggestions": "1. What do theorem 2 and 3 imply intuitively? \n2. Line 230, is the Bellman operator missing?\n3. Why does Eq(1) hold?  \n4. Line 55. The MDP definition is different from the RL community, where $\\gamma$ acts on reward and value functions. It would be better to make the difference more clear.", "questions_for_authors": "See above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nNuGuesFUD", "reviewer": "Reviewer_J3Z7", "summary": "This paper introduces a flow-matching and diffusion-based generative modeling framework to learn accurate Geometric Horizon Models by integrating temporal difference structure in the learning objective and sampling process. The proposed TD$^2$-DD and TD$^2$-CFM achieve low modeling errors in both success measure and value function, and remain robust in long effective horizons. Furthermore, the TD-based GHM approaches demonstrate strong planning performance with Generalized Policy Improvement, greatly outperforming the baseline based on Forward-Backward representation.\n\n## Update after rebuttal\nI thank the authors for their detailed response. Most of my concerns have been addressed. I will keep my positive score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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).\n\n> **Impact of dataset quality and scale on planning and modeling performance**\n\n\nWe 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.\n\nFor 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.\n\n> **Task diversity and inclusion of discrete/stochastic domains**\n\n\nThank 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.\n\nAs 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.\n\nWe 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.\n\n> **Use of TD and potential training instability due to bootstrapping**\n\n\nThank 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.\"\n\nThat 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.\n\nIn 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.\n\n> **Formatting**\n- **L189**: Thank you—we will correct the formatting issue in the final version.\n\n---\n\nWe hope our responses help to contextualize the design choices and emphasize the broader relevance of our contributions.\n\n### References\n[1] Ahmed Touati and Yann Ollivier. Learning One Representation to Optimize All Rewards. Neural Information Processing Systems (NeurIPS), 2021.\n\n\n[2] Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does Zero-Shot Reinforcement Learning Exist? International Conference on Learning Representations (ICLR), 2023.", "claims_and_evidence": "Please see the “Other Strengths and Weaknesses” and “Questions” sections below", "methods_and_evaluation_criteria": "Please see the “Other Strengths and Weaknesses” and “Questions” sections below", "theoretical_claims": "No critical issues are observed in my review of theoretical claims.", "experimental_designs_or_analyses": "Please see the “Other Strengths and Weaknesses” and “Questions” sections below", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n- The proposed method innovatively builds GHMs on flow and diffusion models, while injecting temporal difference structure into the sampling procedure and learning objective.\n- Extensive evaluations demonstrate the efficacy of TD$^2$-CFM and TD$^2$-DD through non-trivial improvement over generative metrics and planning performance.\n- TD$^2$-CFM remains effective and robust in curved path designs.  \n\nWeaknesses:\n- It would be helpful to show how the quality and scale of offline datasets might affect the planning performance and modeling accuracy\n- FB[1] evaluates tasks of different types, including complex maze navigation, robotic manipulation, and games, whereas the evaluations in Section 5.2 mainly focus on continuous locomotion tasks (3 out of 4 domains), including Walker, Cheetah, and Quadruped. Since TD$^2$-based GHM demonstrates the capabilities of modeling successor measures over a long effective horizon, tasks that require long-horizon reasoning, such as navigation tasks in a large Pointmaze or the four-room setting of MiniGrid, should also be tested.\n- So far the evaluation seems mainly performed in low-dimensional but deterministic environments. Given the strong capabilities of flow and diffusion models, it would be a meaningful investigation to understand their performance in stochastic or partially observable environments.\n\n[1] Touati, A. and Ollivier, Y. Learning one representation to optimize all rewards. In Neural Information Processing Systems (NeurIPS), 2021.", "other_comments_or_suggestions": "1. Please re-format the equation (the parentheses in the denominator) on the left column of Line 189.", "questions_for_authors": "1. Bootstrapping is one cause of the training instability in RL when combined with off-policy learning and function approximation. How would the use of bootstrapped samples affect the training stability of the proposed GHMs overall?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Eg1WVIIJFo", "reviewer": "Reviewer_SPuZ", "summary": "This paper studies the problem of learning generative horizon models (GHMs), which are generative models of the successor measure of a policy. That is, GHMs are models capable of generating samples from the discounted distribution of future states visited by a policy at any given time step $t$. In particular, this paper introduces Temporal Difference Flows (TD-Flow), a class of flow-matching techniques used to learn GHMs. The authors introduce three variants of TD-Flow (TD-CFM, Coupled TD-CFM, TD^2-CFM) and additional extensions to diffusion models (TD-DD, TD^2-DD) that improve the accuracy and stability of the learned GHMs. The methods are evaluated in robotics domains from the DeepMind Control Suite with respect to different error metrics, as well as in terms of performance when the methods are employed in a zero-shot transfer setting in combination with generalized policy improvement (GPI).\n\n---\n\n### Post Rebuttal\n\nI thank the authors for their careful response. I am maintaining my acceptance score since I do not have additional questions.\n\nHowever, I would like to emphasize the importance of including the confidence intervals for the figures in the revised manuscript.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> **Number of seeds and statistical reporting**\n\n\nThank 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.\n\n> **Related work and comparison to GPI techniques**\n\n\nThank 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.\n\n> **Clarifications and presentation improvements**\n\n\n- **Line 29**: Thank you for catching this error, as defined in the appendix this should read geometric not generative.\n- **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.\n- **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.\n- **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.\n\n---\n\nWe’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.\n\n### References\n[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.", "claims_and_evidence": "Overall, the paper provides strong empirical and theoretical support for its claims, with extensive benchmarking and mathematical justification.\n\nRegarding its theoretical results, the paper offers convergence guarantees for the TD-based flow methods, demonstrating that TD2-CFM, TD-CFM, and TD-CFM(C) all converge to a fixed probability path​ (assuming that the flow-matching loss is minimized exactly at each iteration).\n\nRegarding the empirical validation, the results indicate that TD^2-based methods outperform alternative models (GANs, VAEs, and diffusion models) in long-horizon predictions, with significantly lower error metrics such as Earth Mover’s Distance and value-function mean squared error.\n\nThe paper also includes ablation experiments to analyze the effect of path design, confirming the robustness of the introduced techniques.", "methods_and_evaluation_criteria": "The proposed methods are evaluated in 4 domains (Maze, Walker, Cheetah, Quadruped) from the DeepMind Control Suite. This benchmark is standard in the RL literature and provides diverse control challenges for evaluating the proposed method. One concern I have is that these domains are almost deterministic. It would be interesting to see how the proposed methods perform in domains with more stochastic transitions, e.g., with multi-modal distributions of next states.\n\nRegarding the baselines employed in the experiments, they represent current state-of-the-art techniques for learning GHMs and make sense for the empirical evaluation.", "theoretical_claims": "I checked the proof of Theorems 1, 2, and 3 and did not find any issues. I did not check the proofs of other theoretical results in the paper and in Appendix E.", "experimental_designs_or_analyses": "The experimental design of the paper is well-conducted and includes an evaluation of the proposed techniques with different relevant metrics (log-likelihood, Earth Mover’s distance, and value function MSE). The authors also investigate how the performance degrades with larger horizons, and how they affect the performance of an agent that employs generalized policy improvement (GPI) for solving tasks.\n\nMy only concern with the empirical design is that the authors employ only three random seeds for each domain. The authors should justify the use of such a small value. Moreover, the results in Figure 2 and Figure 3 should also report the confidence level or statistical dispersion of the reported metrics.", "relation_to_broader_scientific_literature": "This paper's main contributions extend previous works based on the successor representation/measure (SR), especially regarding methods for learning generative horizon models (GHMs) of the SR. Previous works have explored learning GHMs using techniques such as GANs, Normalizing Flows, and auto-encoders. This paper introduced new algorithms and theoretical results based on flow match and diffusion models for learning GHMs. Importantly, their methods are proven to converge to the successor measure and are shown to reduce the variance of sample-based gradient estimates.\n\nThis paper is also related to policy transfer methods based on generalized policy improvement (GPI). In particular, the authors employ their introduced techniques to perform GPI using GHMs conditioned on policy embeddings.", "essential_references_not_discussed": "To the best of my knowledge, the paper discusses the most relevant current state-of-the-art methods for learning GHMs.", "other_strengths_and_weaknesses": "Besides the points raised in the other sections, one of the paper's strengths is its clarity and presentation. In particular, the paper covers many different mathematical formulations (e.g., RL, flow matching, diffusion models) while maintaining mathematical rigor, clarity, and precision.\n\nOne point for improvement would be the addition of a dedicated related work section that pinpoints the most relevant contributions of the paper in contrast to the state-of-the-art. Currently, these points are scattered throughout the paper.", "other_comments_or_suggestions": "- In line 29, Thakoor et al., 2022 refer to GHM as “geometric horizon models”, and not  “generative horizon models”.\n\n- In line 59, why restrict to deterministic policies?\n\n- In Equations 2 and 3, the variable $\\rho$ was not defined.\n\n- In Section 5.2, I suggest discussing in more depth the relation with the GPI technique of (Thakoor\net al., 2022), which also extends GPI via GHMs. Currently, in this section, the authors only mention the work that first introduced GPI (Barreto et al., 2017). Did the authors try to perform GPI similarly to Thakoor et al, 2022, that is, using geometric switching policies?\n\n- The introduced method had many variants (TD-CFM, Coupled TD-CFM, TD^2-CFM) and additional extensions to diffusion models (TD-DD, TD^2-DD). This makes it difficult for the reader to track the disadvantages and advantages of each variant, as well as their unique properties. I suggest the authors include a table that summarizes all introduced variants with their properties and pros and cons for improved clarity of the paper.", "questions_for_authors": "See Other Comments Or Suggestions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper leverages a novel Bellman equation structure along the probability paths to improve the learning of geometric horizon models, significantly enlarging their prediction horizon. Theoretical analysis is provided accompanied with convincing empirical results. All the reviewers appreciate the solid contribution and the clear presentation. I, therefore, recommend accept.
Decision: Accept (oral)