Inference via Interpolation: Contrastive Representations Provably Enable Planning and Inference

Authors: benjamin  eysenbach
vivek  myers
ruslan  salakhutdinov
sergey  levine
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 02:11:40

Abstract: Given time series data, how can we answer questions like ``what will happen in the future?'' and ``how did we get here?'' These sorts of probabilistic inference questions are challenging when observations are high-dimensional. In this paper, we show how these questions can have compact, closed form solutions in terms of learned representations. The key idea is to apply a variant of contrastive learning to time series data. Prior work already shows that the representations learned by contrastive learning encode a probability ratio. By extending prior work to show that the marginal distribution over representations is Gaussian, we can then prove that joint distribution of representations is also Gaussian. Taken together, these results show that representations learned via temporal contrastive learning follow a Gauss-Markov chain, a graphical model where inference (e.g., prediction, planning) over representations corresponds to inverting a low-dimensional matrix. In one special case, inferring intermediate representations will be equivalent to interpolating between the learned representations. We validate our theory using numerical simulations on tasks up to 46-dimensions.

Main paper: 
 introduction: 
 Probabilistic modeling of time-series data has applications ranging from robotic control [84] to material science [42] , from cell biology [74] to astrophysics [55] . These applications are often concerned with two questions: predicting future states (e.g., what will this cell look like in an hour), and inferring trajectories between two given states. However, answering these questions often requires reasoning over high-dimensional data, which can be challenging as most tools in the standard probabilistic toolkit require generation. Might it be possible to use discriminative methods (e.g., contrastive learning) to perform such inferences? Many prior works aim to learn representations that are easy to predict while retaining salient bits of information. For time-series data, we want the representation to remain a sufficient statistic for distributions related to time -for example, they should retain bits required to predict future states (or representations thereof). While generative methods [26, 56, 101, 103] have this property, they tend to be computationally expensive (see, e.g., [72] ) and can be challenging to scale to high-dimensional observations. ψ A Latent Space • small dim • easy planning Observation Space • large dim • hard planning Figure 1 : We apply temporal contrastive learning to observation pairs to obtain representations (ψ(x0), ψ(x t+k )) such that Aψ(x0) is close to ψ(x t+k ). While inferring waypoints in the high-dimensional observation space is challenging, we show that the distribution over intermediate latent representations has a closed form solution corresponding to linear interpolation between the initial and final representations.In this paper, we will study how contrastive methods (which are discriminative, rather than generative) can perform inference over times series. Ideally, we want representations of observations x to be a sufficient statistic for temporal relationships (e.g., does x ′ occur after x?) but need not retain other information about x (e.g. the location of static objects). This intuition motivates us to study how contrastive representation learning methods [18, 65, 81, 87, 96] might be used to solve prediction and planning problems on time series data. While prior works in computer vision [16, 65] and natural language processing (NLP) [59] often study the geometry of learned representations, our results show how geometric operations such as interpolation are related to inference. Our analysis will focus on a regularized version of the symmetrized infoNCE objective [71] , generating positive examples by sampling pairs of observations from the same time series data. We will study how representations learned in this way can facilitate two inference questions: prediction and planning. 2 As a stepping stone, we will build upon prior work [90] to show that regularized contrastive learning should produce representations whose marginal distribution is an isotropic Gaussian distribution.The main contribution of this paper is to demonstrate how intermediate and future time steps in a time series can be inferred easily using contrastive representations. This inference problem captures a number of practical tasks: interpolation, in-filling, and even planning and control, where the intermediate steps represent states between a stand and goal. While ordinarily these problems require an iterative inference or optimization procedure, with contrastive representations this can be done simply by inverting a low-dimensional matrix. In one special case, inference will correspond to linear interpolation. Our first step is to prove that, under certain assumptions, the distribution over future representations has a Gaussian distribution, with a mean that is a linear function of the initial state representation (Lemma 1). This paves the main to our main result (Theorem 4.1): given an initial and final state, we show that the posterior distribution over an intermediate state representations also follows a Gaussian distribution. Said in other words, the representations follow a Gauss-Markov chain, 3 wherein any joint or conditional distribution can be computed by inverting a low-dimensional matrix [57, 93] (See Fig. 1 ). In one special case, inference will correspond to linearly interpolating between the representations of an initial state and final state. Section 5 provides numerical experiments. 
related work: 
 Representations for time-series data. In applications ranging from robotics to vision to NLP, users often want to learn representations of observations from time series data such that the spatial arrangement of representations reflects the temporal arrangement of the observations [28, 59, 65, 70] . Ideally, these representations should retain information required to predict future observations and infer likely paths between pairs of observations. Many approaches use an autoencoder, learning rep-resentations that retain the bits necessary to reconstruct the input observation, while also regularizing the representations to compressed or predictable [12, 20, 43, 45, 68, 104] . A prototypical method is the sequential VAE [101] , which is computationally expensive to train because of the reconstruction loss, but is easy to use for inference. Our work shares the aims of prior prior methods that attempt to linearize the dynamics of nonlinear systems [7, 23, 63, 64, 80, 92] , including videos [33, 40] . Our work aims to retain uncertainty estimates over predictions (like the sequential VAE) without requiring reconstruction. Avoiding reconstruction is appealing practically because it decreases the computational requirements and number of hyperparameters; and theoretically because it means that representations only need to retain bits about temporal relationships and not about the bits required to reconstruct the original observation.Contrastive Learning. Contrastive learning methods circumvent reconstruction by learning representations that merely classify if two events were sampled from the same joint distribution [17, 35, 71] . When applied to representing states along trajectories, contrastive representations learn to classify whether two points lie on the same trajectory or not [29, 65, 70, 78, 97] . Empirically, prior work in computer vision and NLP has observed that contrastive learning acquires representations where interpolation between representations corresponds to changing the images in semantically meaningful ways [19, 50, 59, 67, 95, 98] . Our analysis will be structurally similar to prior theoretical analysis on explaining why word embeddings can solve analogies [3, 5, 48] . Our work will make a Gaussianity assumption similar to Arora et al. [5] and our Markov assumption is similar to the random walks analyzed in Arora et al. [5] , Hashimoto et al. [36] . Our paper builds upon and extends these results to answer questions such as: "what is the distribution over future observations representations?" and "what is the distribution over state (representations) that would occur on the path between one observation and another?" While prior work is primarily aimed at explaining the good performance of contrastive word embeddings (see, e.g., [5] ), we are primarily interested in showing how similar contrastive methods are an effective tool for inference over high-dimensional time series data. Our analysis will show how representations learned via temporal contrastive learning (i.e., without reconstruction) are sufficient statistics for inferring future outcomes and can be used for performing inference on a graphical model (a problem typically associated with generative methods).Goal-oriented decision making. Much work on time series representations is done in service of learning goal-reaching behavior, an old problem [46, 61] that has received renewed attention in recent years [14, 15, 21, 37, 39, 52, 76, 99] . Some of the excitement in goal-conditioned RL is a reflection of the recent success of self-supervised methods in computer vision [73] and NLP [66] . Our analysis will study a variant of contrastive representation learning proposed in prior work for goal-conditioned RL [29, 78] . These methods are widespread, appearing as learning objectives for learning value functions [1, 28, 29, 49, 51, 53, 60, 86, 91, 102] , as auxiliary objectives [4, 9, 13, 60, 77, 82, 83] , in objectives for model-based RL [2, 32, 58, 80] , and in exploration methods [25, 34] . Our analysis will highlight connections between these prior methods, the classic successor representation [8, 24] , and probabilistic inference.Planning. Planning lies at the core of many RL and control methods, allowing methods to infer the sequence of states and actions that would occur if the agent navigated from one state to a goal state. While common methods such as PRM [44] and RRT [47] focus on building random graphs, there is a strong community focusing on planning methods based on probabilistic inference [6, 85, 94] .The key challenge is scaling to high-dimensional settings. While semi-parametric methods make progress on this problem this limitation through semi-parametric planning [27, 30, 100] , it remains unclear how to scale any of these methods to high-dimensional settings when states do not lie on a low-dimensional manifold. Our analysis will show how contrastive representations may lift this limitation, with experiments validating this theory on 39-dimensional and 46-dimensional tasks. 
preliminaries: 
 Our aim is to learn representations of time series data such that the spatial arrangement of representations corresponds to the temporal arrangement of the underlying data: if one example occurs shortly after another, then they should be mapped to similar representations. This problem setting arises in many areas, including video understanding and reinforcement learning. To define this problem formally, we will define a Markov process with states x t indexed by time t: 4 p(x 1:T | x 0 ) = T t=0 p(x t+1 | x t ). The dynamics p(x t+1 | x t ) tell us the immediate next state, and we can define the distribution over states t steps in the future by marginalizing over the intermediate states, p t (x t | x 0 ) = p(x 1:t | x 0 ) dx 1:t-1 . A key quantity of interest will be the γ-discounted state occupancy measure, which corresponds to a time-averaged distribution over future states:EQUATIONContrastive learning. Our analysis will focus on applying contrastive learning to a particular data distribution. Contrastive learning [35, 65, 75] acquires representations using "positive" pairs (x, x + ) and "negative" pairs (x, x -). While contrastive learning typically learns just one representation, we will use two different representation for the two elements of the pair; that is, our analysis will use terms like ϕ(x), ψ(x + ) and ψ(x -). We assume all representations lie in R k .The aim of contrastive learning is to learn representations such that positive pairs have similar representations (ϕ(x) ≈ ψ(x + )) while negative pairs have dissimilar representations (ϕ(x) ̸ = ψ(x -)). Let p(x, x + ) be the joint distribution over positive pairs (i.e., (x, x + ) ∼ p(x, x + )). We will use the product of the marginal distributions to sample negative pairs ((x, x -) ∼ p(x)p(x)). Let B be the batch size, and note that the positive samples x + j at index j in the batch serve as negatives for x i for any i ̸ = j. Our analysis is based on the infoNCE objective without resubstitution [65, 81] :max ϕ(•),ψ(•) E {(xi,x + i )} B i=1 ∼p(x,x + ) B i=1 log e -1 2 ∥ϕ(x i )-ψ(x + i )∥ 2 2 j̸ =i e -1 2 ∥ϕ(x i )-ψ(x + j )∥ 2 2 + log e -1 2 ∥ϕ(x i )-ψ(x + i )∥ 2 2 j̸ =i e -1 2 ∥ϕ(x j )-ψ(x + i )∥ 2 2(2)We will use the symmetrized version of this objective [71] , where the denominator is the sum across rows of a logits matrix and once where it is a sum across the.While contrastive learning is typically applied to an example x and an augmentation x + ∼ p(x | x) of that same example (e.g., a random crop), we will follow prior work [65, 78] in using the time series dynamics to generate the positive pairs, so x + will be an observation that occurs temporally after x. While our experiments will sample positive examples from the discounted state occupancy measure (x + ∼ p t+ (x t+ | x)) in line with prior work [29] , our analysis will also apply to different distributions (e.g., always sampling a state k steps ahead).While prior work typically constrains the representations to have a constant norm (i.e., to lie on the unit hypersphere) [65] , we will instead constrain the expected norm of the representations is bounded, a difference that will be important for our analysis:EQUATIONBecause the norm scales with the dimension of the representation, we have scaled down the left side by the representation dimension, k. In practice, we will impose this constraint by adding a regularization term λ E p(x) ∥ψ(x)∥ 2 2 to the infoNCE objective (Eq. 2) and dynamically tuning the weight λ via dual gradient descent.This section outlines the two key assumptions behind our analysis, both of which have some theoretical justification. Our main assumption examines the distribution over representations: Assumption 1. Regularized, temporal contrastive learning acquires representations whose marginal distribution representations p(ψ) ≜ p(x) 1(ψ(x) = ψ) dx is an isotropic Gaussian distribution:EQUATIONIn Appendix A.1 we extend prior work [90] provide some theoretical intuition for why this assumption should hold: namely, that the isotropic Gaussian is the distribution that maximizes entropy subject to an expected L2 norm constraint (Eq. 3) [22, 41, 79] . Our analysis also assumes that the learned representations converge to the theoretical minimizer of the infoNCE objective: Assumption 2. Applying contrastive learning to the symmetrized infoNCE objective results in representations that encode a probability ratio:EQUATIONThis assumption holds under ideal conditions [54, 69] (see Appendix A.5), 5 but we nonetheless call this an "assumption" because it may not hold in practice due to sampling and function approximation error. This assumption means the learned representations are sufficient statistics for predicting the probability (ratio) of future states: these representations must retain all the information pertinent to reasoning about temporal relationships, but need not retain information about the precise contents of the observations. As such, they may be much more compressed than representations learned via reconstruction.Combined, these assumptions will allow us to express the distribution over sequences of representations as a Gauss-Markov chain. The denominator in Assumption 2, p(x), may have a complex distribution, but Assumption 1 tells us that the distribution over representations has a simpler form. This will allow us to rearrange Assumption 2 to express the conditional distribution over representations as the product of two Gaussian likelihoods. Note that the left hand side of Assumption 2 already looks like a Gaussian likelihood. 
contrastive representations make inference easy: 
 In this section, our main result will be to show how representations learned by (regularized) contrastive learning are distributed according to a Gauss-Markov chain, making it straightforward to perform inference (e.g., planning, prediction) over these representations. Our proof technique will combine (known) results about Gaussian distributions with (known) results about contrastive learning. We start by discussing an important choice of parametrization (Section 4.1) that facilitates prediction (Section 4.2) before presenting the main result in Section 4.3. This section describes the two encoders (ψ(•), ϕ(•)) to compute representations of x and x + . While prior work in computer vision and NLP literature use the same encoder for both x and x + , this decision does not make sense for many time-series data as it would imply that our prediction for p(x t | x 0 ) is the same as our prediction for p(x 0 | x t ). However, the difficulty of transiting from x 0 to x t (e.g., climbing to the peak of a mountain) might be more difficult than the reverse (e.g., sledding down a mountain). Our proposed parametrization will handle this asymmetry.ψ(x 0 ) ϕ(x 0 ) = Aψ(x 0 ) A ψ(x t )We will treat the encoder ψ(•) as encoding the contents of the state. We will additionally learn a matrix A so that the function ψ → Aψ corresponds to a (multi-step) prediction of the future representation. To map this onto contrastive learning, we will use ϕ(x) ≜ Aψ(x) as the encoder for the initial state. One way of interpreting this encoder is as an additional linear projection applied on top of ψ(•), a design similar to those used in other areas of contrastive learning [17] . Once learned, we can use these encoders to answer questions about prediction (Section 4.2) and planning (Section 4.3). Given an initial state x 0 , what states are likely to occur in the future? Answering this question directly in terms of high-dimensional states is challenging, but our learned representations provide a straightforward answer. Let ψ 0 = ψ(x 0 ) and ψ t+ = ψ(x t+ ) be random variables representing the representations of the initial state and a future state. Our aim is to estimate the distribution over these future representations, p(ψ t+ | ψ 0 ). We will show that the learned representations encode this distribution. Lemma 1. Under the assumptions from Section 3, the distribution over representations of future states follows a Gaussian distribution with mean parameter given by the initial state representation:ψ(x 0 ) Aψ(x 0 ) c c+1 Aψ(x 0 )EQUATIONThe main takeaway here is that the distribution over future representations has a convenient, closed form solution. The representation norm constraint, c, determines the shrinkage factor c c+1 ∈ [0, 1); highly regularized settings (small c) move the mean closer towards the origin and decrease the variance, as visualized in Fig. 3 . Regardless of the constraint c, the predicted mean is a linear function ψ → c c+1 Aψ. The proof is in Appendix A.2. The proof technique is similar to that of the law of the unconscious statistician.We now show how these representations can be used for a specific type of planning: given an initial state x 0 and a future state x t+ , infer the representation of an intermediate "waypoint" state x w . The next section will extend this analysis to inferring the entire sequence of intermediate states. We assume x 0 → x w → x t+ form a Markov chain where x w ∼ p(x t+ | x 0 = x 0 ) and x t+ ∼ p(x t+ | x 0 = x w ) are both drawn from the discounted state occupancy measure (Eq. 1). Let random variable ψ w = ψ(x w ) be the representation of this intermediate state. Our main result is that the posterior distribution over waypoint representations has a closed form solution in terms of the initial state representation and future state representation: Theorem 4.1. Under Assumptions 1 and 2, the posterior distribution over waypoint representations is a Gaussian whose mean and covariance are linear functions of the initial and final state representations:p(ψ w | ψ 0 , ψ t+ ) = N ψ w ; µ = Σ(A T ψ t+ + Aψ 0 ), Σ -1 = c c+1 A T A + c+1 c I .The proof (Appendix A.3) uses the Markov property together with Lemma 1. The main takeaway from this lemma is that the posterior distribution takes the form of a simple probability distribution (a Gaussian) with parameters that are linear functions of the initial and final representations.Example 1: A = I and the c is very large (little regularization). Then, the covariance is Σ -1 ≈ 2I and the mean is the simple average of the initial and final representations µ ≈ 1 2 (ψ 0 + ψ t+ ). In other words, the waypoint representation is the midpoint of the line ψ 0 → ψ t+ .Example 2: A is a rotation matrix and c is very large. Rotation matrices satisfy A T = A -1 so the covariance is again Σ -1 ≈ 2I. As noted in Section 4.2, we can interpret Aψ 0 as a prediction of which representations will occur after ψ 0 . Similarly, A -1 ψ t+ = A T ψ t+ is a prediction of which representations will occur before ψ t+ . Theorem 4.1 tells us that the mean of the waypoint distribution is the simple average of these two predictions, µ ≈ 1 2 (A T ψ t+ + Aψ 0 ). Example 3: A is a rotation matrix and c = 0.01 (very strong regularization). In this case Σ -1 = 0.01 0.01+1 A T A+ 0.01+1 0.01 I ≈ 100I, so µ ≈ 1 100 (ψ 0 +ψ t+ ) ≈ 0. Thus, in the case of strong regularization, the posterior concentrates around the origin.This section extends the analysis to multiple intermediate states. Again, we will infer the posterior distribution of the representations of these intermediate states, ψ w1 , ψ w2 , • • • . We assume that these states form a Markov chain. Theorem 4.2. Given observations from a Markov chain x 0 → x 1 • • • x t+ , the joint distribution over representations is a Gaussian distribution. Using ψ 1:n = (ψ w1 , • • • , ψ wn ) to denote the concatenated representations of each observation, we can write this distribution asp(ψ 1:n ) ∝ exp -1 2 ψ T 1:n Σ -1 ψ 1:n + η T ψ 1:n ,where Σ -1 is a tridiagonal matrixΣ -1 = c c+1 A T A+ c+1 c I -A T -A c c+1 A T A+ c+1 c I -A T . . . and η =   Aψ0 0 . . . A T ψt+   .This distribution can be written in the canonical parametrization as Σ = Λ -1 and µ = Ση. Recall that Gaussian distributions are closed under marginalization. Thus, once in this canonical parametrization, the marginal distributions can be obtained by reading off individual entries of these parameters:p(ψ i | ψ 0 , ψ t+ ) = N ψ i ; µ i = (Ση) (i) , Σ i = (Λ -1 ) (i,i) .The key takeaway here is that this posterior distribution over waypoints is Gaussian, and it has a closed form expression in terms of the initial and final representations (as well as regularization parameter c and the learned matrix A).In the general case of n intermediate states, the posterior distribution isp(ψ w1 • • • ψ wn | ψ 0 , ψ t+ ) ∝ e -1+ 1 c 2 n i=1 ∥ c c+1 Aψw i -ψw i+1 ∥ 2 2, where ψ w0 = ψ 0 and ψ wn+1 = ψ t+ . This corresponds to a chain graphical model with edge potentialsf (ψ, ψ ′ ) = e -1+ 1 c 2 ∥ c c+1 Aψ-ψ ′ ∥ 22 . Special case. To build intuition, consider the special case where A is a rotation matrix and c is very large, so c c+1 A T A + c+1 c ≈ 2I. In this case, Σ -1 is a (block) second difference matrix [38] :Σ -1 = 2I -I -I 2I -I . . . .The inverse of this matrix has a closed form solution [62, Pg. 471], allowing us to obtain the mean of each waypoint in closed form:EQUATIONwhere λ(i) = i n+1 . Thus, each posterior mean is a convex combination of the (forward prediction from the) initial representation and the (backwards prediction from the) final representation. When A is the identity matrix, the posterior mean is simple linear interpolation between the initial and final representations! 
numerical simulation: 
 We include several didactic experiments to illustrate our results. All results and figures can be reproduced by running make in the source code: https://github.com/vivekmyers/contrastive_ planning. The expected compute time is a few hours on a A6000 GPU. Figures in this section show error across different training and dataset split seeds.To validate our analysis, we design a time series task with 2D points where inference over intermediate points (i.e., in-filling) requires nonlinear interpolation. Fig. 4 (Top Left) shows the dataset of time series data, starting at the origin and spiraling outwards, with each trajectory using a randomly-chosen initial angle. We applied contrastive learning with the parametrization in Section 4.1 to these data and used the learned representations to solve prediction and planning problems (see Fig. 4 for details).Note that these predictions correctly handle the nonlinear structure of these data -states nearby the initial state in Euclidean space that are not temporally adjacent are assigned low likelihood.Our next experiment studies whether the inferred representations are useful for solving a control task. We took a 2d maze environment and dataset from prior work (Fig. 5 , Left) [31] and learned encoders from this dataset. To solve the maze, we take the observation of the starting state and goal state, compute the representations of these states, and use the analysis in Section 4.3 to infer the sequence of intermediate representations. We visualize the results using a nearest neighbor retrieval (Fig. 5 , Left). Appendix Fig. 7 contains additional examples. Finally, we studied whether these representations are useful for control. We implemented a simple proportional controller for this maze. As expected, this proportional controller can successfully navigate to close goals, but fails to reach distant goals (Fig. 5 , Right). However, if we use the proportional controller to track a series of waypoints planned using our representations (i.e., the orange dots shown in Fig. 5 (Left)), the success rate increases by up to 4.5×. To test the importance of nonlinear representations, we compare with a "PCA" baseline that predicts waypoints by interpolating between the principal components of the initial state and goal state. The better performance of our method indicates the importance of doing the interpolation using representations that are nonlinear functions of the input observations. While prior methods learn representations to encode temporal distances, it is unclear whether these methods support inference via interpolation. To test this hypothesis, we use one of these methods ("VIP" [51] ) as a baseline. While the VIP representations likely encode similar bits as our representations, the better performance of the contrastive representations indicates that the VIP representations do not expose those bits in a way that makes planning easy.In this section we provide preliminary experiments showing the planning approach in Section 4 scales to higher dimensional tasks. We used two datasets from prior work [31] : door-human-v0 (39-dimensional observations) and hammer-human-v0 (46-dimensional observations). After learning encoders on these tasks, we evaluated the inference capabilities of the learned representations. Given the first and last observation from a trajectory in a validation set, we use linear interpolation (see Eq. ( 7)) to infer the representation of five intermediate waypoint representations.We evaluate performance in two ways. Quantitatively, we measure the mean squared error between each of the true waypoint observations and those inferred by our method. Since our method infers representations, rather than observations, we use a nearest-neighbor retrieval on a validation set so that We use a dataset of trajectories demonstrating door opening from prior work [31] to learn representations. (Top Right) We use our method and three baselines to infer one intermediate waypoint between the first and last observation in a trajectory from a held-out validation set. Errors are measured using the mean squared error with the true waypoint observation; predicted representations are converted to observations using nearest neighbors on a validation set. (Bottom) We visualize a TSNE [89] of the states along the sampled trajectory as blue circles, with the transparency indicating the index along the trajectory. The inferred plan is shown as red circles connected by arrows. Our method generates better plans than alternative representation learning methods (PCA, VIP).we can measure errors in the space of observations. Qualitatively, we visualize the high-dimensional observations from the validation trajectory using a 2-dimensional TSNE [89] embedding, overlying the infer waypoints from our method; as before, we convert the representations inferred by our method to observations using nearest neighbors.We compare with three alternative methods in Fig. 6 . To test the importance of representation learning, we first naïvely interpolate between the initial and final observations ("no planning"). The poor performance of this baseline indicates that the input time series are highly nonlinear. Similarly, interpolating the principle components of the initial and final observations ("PCA") performs poorly, again highlighting that the input time series is highly nonlinear and that our representations are doing more than denoising (i.e., discarding directions of small variation). The third baseline, "VIP" [51] , learns representations to encode temporal distances using approximate dynamic programming. Like our method, VIP avoids reconstruction and learns nonlinear representations of the observations. However, the results in Fig. 6 highlight that VIP's representations do not allow users to plan by interpolation. The error bars shown in Fig. 6 (Top Right) show the standard deviation over 500 trajectories sampled from the validation set. For reproducibility, we repeated this entire experiment on another task, the 46-dimensional hammer-human-v0 from D4RL. The results, shown in Appendix Fig. 8 , support the conclusions above. Taken together, these results show that our procedure for interpolating contrastive representations continues to be effective on tasks where observations have dozens of dimensions. 
discussion: 
 Representation learning is at the core of many high-dimensional time-series modeling questions, yet how those representations are learned is often disconnected with the inferential task. The main contribution of this paper is to show how discriminative techniques can be used to acquire compact representations that make it easy to answer inferential questions about time. The precise objective and parametrization we studied is not much different from that used in practice, suggesting that either our theoretical results might be adapted to the existing methods, or that practitioners might adopt these details so they can use the closed-form solutions to inference questions. Our work may also have implications for studying the structure of learned representations. While prior work often studies the geometry of representations as a post-hoc check, our analysis provides tools for studying when interpolation properties are guaranteed to emerge, as well as how to learn representations with certain desired geometric properties.Limitations. Our analysis hinges on the two assumptions mentioned in Section 3.1, and it remains open how errors in those approximations translate into errors in our analysis. One important open question is whether it is always possible to satisfy these assumptions using sufficiently-expressive representations.The infoNCE objective (Eq. ( 2)) can be decomposed into an alignment term and a uniformity term [90] , where the uniformity term can be simplified as follows:E x∼p(x) log E x -∼p(x) e -1 2 ∥Aψ(x -)-ψ(x)∥ 2 2 = 1 N N i=1 log 1 N -1 j=1•••N,j̸ =i e -1 2 ∥Aψ(xi)-ψ(xj )∥ 2 2 = 1 N N i=1 log 1 N -1 j=1•••N,j̸ =i 1 (2π) k/2 e -1 2 ∥Aψ(xi)-ψ(xj )∥ 2 2 N (µ=ψ(xj );Σ=I) + k 2 log(2π) = 1 N N i=1 log pGMM (ψ(x i )) + k 2 log(2π) = -Ĥ[ψ(x)] + k 2 log(2π).The derivation above extends that in Wang and Isola [90] by considering a Gaussian distribution rather than a von Mises Fisher distribution. We are implicitly making the assumption that the marginal distributions satisfy p(x) = p(x -). This difference corresponds to our choice of using a negative squared L2 distance in the infoNCE loss rather than an inner product, a difference that will be important later in our analysis. A second difference is that we do not use the resubstitution estimator (i.e., we exclude data point x i from our estimate of pGMM when evaluating the likelihood of x i ), which we found hurt performance empirically. The takeaway from this identity is that maximizing the uniformity term corresponds to maximizing (an estimate of) the entropy of the representations.We next prove that the maximum entropy distribution with an expected L2 norm constraint is a Gaussian distribution. Variants of this result are well known [22, 41, 79] , but we include a full proof here for transparency.Lemma 2. The maximum entropy distribution satisfying the expected L2 norm constraint in Eq. ( 3) is a multivariate Gaussian distribution with mean µ = 0 and covariance Σ = c • I Proof. We start by defining the corresponding Lagrangian, with the second constraint saying that p(x) must be a valid probability distribution.L(p) = H p [x] + λ 1 E p(x) ∥x∥ 2 2 -c • k + λ 2 p(x) dx -1We next take the derivative w.r.t. p(x):∂L ∂p(x) = -p(x)/p(x) -log p(x) + λ 1 ∥x∥ 2 2 + λ 2Setting this derivative equal to 0 and solving for p(x), we getp(x) = e -1+λ2+λ1∥x∥ 2 2 .We next solve for λ 1 and λ 2 to satisfy the constraints in the Lagrangian. Note that x ∼ N (µ = 0, Σ = c • I) has an expected norm E[∥x∥ 2 2 ] = c • k, so we must have λ 1 = -1 2c . We determine λ 1 as the normalizing constant for a Gaussian, finally giving us:p(x) = 1 (2cπ) k/2 e -1 2c ∥x∥ 2 2corresponding to an isotropic Gaussian distribution with mean µ = 0 and covariance Σ = c • I.Below we present the proof of Lemma 1 Proof. Our proof technique will be similar to that of the law of the unconscious statistician:p(ψ t+ | ψ 0 ) (a) = p(ψt+,ψ0) p(ψ0) ∝ p(ψ t+ , x t+ , ψ 0 , x 0 ) dx t+ dx 0 (b) = p(ψ t+ | x t+ )p(ψ 0 | x 0 )p(x t+ | x 0 )p(x 0 ) dx t+ dx 0 (c) ∝ 1(ψ(x t+ ) = ψ t+ ) 1(ψ(x 0 ) = ψ 0 )p(x t+ )e -1 2 ∥Aψ(x0)-ψ(xt+)∥ 2 2 p(x 0 ) dx t+ dx 0 (d) = e -1 2 ∥Aψ0-ψt+∥ 2 2 1(ψ(x t+ ) = ψ t+ ) 1(ψ(x 0 ) = ψ 0 )p(x t+ )p(x 0 ) dx t+ dx 0 (e) = e -1 2 ∥Aψ0-ψt+∥ 2 2 p(x t+ ) 1(ψ(x t+ )) dx t+ p(ψt+) p(x 0 ) 1(ψ(x 0 ) dx 0 p(ψ0) (f ) ∝ e -1 2 ∥Aψ0-ψt+∥ 2 2 e -1 2c ∥ψt+∥ 2 2 e -1 2c ∥ψ0∥ 2 2 (g) ∝ e - 1+ 1 c 2 1 1+ 1 c Aψ0-ψt+ 2 2 ∝ N ψ t+ ; µ = c c+1 Aψ 0 , Σ = c c+1 I .In (a) we applied Bayes' Rule and removed the denominator, which is a constant w.r.t. ψ t+ . In (b) we factored the joint distribution, noting that ψ t+ and ψ 0 are deterministic functions of x t+ and x 0 respectively, so they are conditionally independent from the other random variables. In (c) we used Assumption 2 after solving for p(x t+ | x 0 ) = p(x t+ )e -1 2 ∥Aψ(x0)-ψ(x)∥ 2 2. In (d) we noted that when the integrand is nonzero, it takes on a constant value of e -1 2 ∥Aψ0-ψt+∥ 2 2 , so we can move that constant outside the integral. In (e) we used the definition of the marginal representation distribution (Eq. 6). In (f) we used Assumption 1 to write the marginal distributions p(ψ t+ ) and p(ψ 0 ) as Gaussian distributions. We removed the normalizing constants, which are independent of ψ t+ . In (g) we completed the square and then recognized the expression as the density of a multivariate Gaussian distribution.Proof.p(ψ w | ψ 0 , ψ t+ ) (a) = p(ψ t+ | ψ w )p(ψ w | ψ 0 ) p(ψ t+ | ψ 0 ) (b) ∝ e -1+ 1 c 2 ∥ c c+1 Aψw-ψt+∥ 2 2 e -1+ 1 c 2 ∥ c c+1 Aψ0-ψw∥ 2 2 (c) ∝ e -1 2 (ψw-µ) T Σ -1 (ψw-µ) = N (ψ w ; µ, Σ)whereΣ -1 = c c+1 A T A + c+1 c I and µ = Σ(A T ψ t+ + Aψ 0 ).In line (a) we used the definition of the conditional distribution and then simplified the numerator using the Markov property. Line (b) uses the Lemma 1. Line (c) completes the square, the details of which are below:EQUATIONT 0 c c + 1 A T c c + 1 A ψ 0 -2ψ T 0 c c + 1 A T ψ w + ψ T w ψ w const. = 1 2 c + 1 c ψ T w c c + 1 2 A T A + I ψ w -2 c c + 1 A T ψ t+ + Aψ 0 T ψ w = 1 2 ψ T w c c + 1 A T A + c + 1 c I Σ -1 ψ w -A T ψ t+ + Aψ 0 T ψ w const. = (ψ w -µ) T Σ -1 (ψ w -µ),where Σ -1 = c c+1 A T A + c+1 c I and µ = Σ(A T ψ t+ + Aψ 0 ). Above, we have used const.= to denote equality up to an additive constant that is independent of ψ w .Proof. We start by recalling that the waypoints form a Markov chain, so we can express their joint density as a product of conditional densities:p(ψ 1:n ) = p(ψ 0 )p(ψ 1 | ψ 0 )p(ψ 2 | ψ 1 ) • • • p(ψ n | ψ n-1 ).The aim of this lemma is to express the joint distribution over multiple waypoints, given an initial and final state representation:EQUATION= exp -( ( ( ( ( ( (1 2 c c+1 ψ T 0 A T Aψ 0 + ψ T 0 A T ψ 1 -1 2 c+1 c ψ T 1 ψ 1 -1 2 c c+1 ψ T 1 A T Aψ 1 + ψ T 1 A T ψ 2 -1 2 c+1 c ψ T 2 ψ 2 -1 2 c c+1 ψ T 2 A T Aψ 2 + ψ T 2 A T ψ 3 -1 2 c+1 c ψ T 3 ψ 3 . . . -1 2 c c+1 ψ T n A T Aψ n + ψ T n A T ψ t+ - 1 2 c+1 c ψ T t+ ψ t+ = exp -1 2 ψ T 1:n Σ -1 ψ 1:n + η T ψ 1:n , where Σ -1 is a tridiagonal matrix Σ -1 =       c c+1 A T A + c+1 c I -A T -A c c+1 A T A + c+1 c I -A T -A c c+1 A T A + c+1 c I . . .       and η =       Aψ 0 0 . . . 0 A T ψ t+       .In (a) we applied Bayes' rule and removed the denominator because it is a constant with respect to ψ 1:n . In (b) we applied the Markov assumption. In (c) we used Lemma 1 to express the conditional probabilities as Gaussians, ignoring the proportionality constants (which are independent of ψ. In (d) we simplified the exponents, removing terms that do not depend on ψ 1:nAssumption 2 states e -1 when Eq. ( 2) is optimized.We can justify this assumption by analyzing the general solution to the symmetrized version of the Oord et al. [65] infoNCE objective, which we do in Lemma 3. Applying this lemma to our representation learning objective (2) for sufficiently large batch size B then yields Eq. ( 5), with the function approximator ∥ϕ(x) -ψ(x + )∥ 2 ≈ f (x, x + ). Lemma 3. The solution to the optimization problemEQUATIONsatisfiesEQUATIONfor some C.Proof of Lemma 3. We first break down the LHS and RHS of Eq. ( 2):max f lim B→∞ E {(xi,x + i )} B i=1 ∼p(x,x + ) 1 B B i=1 log e f (x i ,x + i ) j̸ =i e f (x i ,x + j ) J1 + log e f (x i ,x + i ) j̸ =i e f (x j ,x + i ) J2 J 1 (f ) = lim B→∞ E {(xi,x + i )} B i=1 ∼p(x,x + ) 1 B B i=1 log e f (x i ,x + i ) j̸ =i e f (x i ,x + j ) J 2 (f ) = lim B→∞ E {(xi,x + i )} B i=1 ∼p(x,x + ) 1 B B i=1 log e f (x i ,x + i ) j̸ =i e f (x j ,x + i )We now use the following result from Ma and Collins [54] :Lemma 4. The optimal solutions f 1 and f 2 for J 1 and J 2 satisfyEQUATIONEQUATIONfor arbitrary c 1 (x), c 2 (x + ).For any C, when c 1 (x) = Cp(x) and c 2 (x + ) = Cp(x + ),EQUATIONIt follows that Eq. ( 12) maximizes both J 1 and J 2 , and is precisely the optimal solution Eq. ( 9) for Eq. ( 8).What does C represent? From Eq. ( 9), we can connect C to the mutual information I(x, x + ):EQUATION)Proof of Lemma 4. We can first consider J 1 without loss of generality. Denoting g(x, x + ) = e f (x,x + ) , we take the functional derivative: = limδJ 1 (log g) = lim B→∞ δ E {(xi,x + i )} B i=1 ∼p(x,x + ) 1 B B i=1 log g(x i , x + i ) j̸ =i g(x i , x + j ) = lim B→∞ E {(xi,x + i )} B i=1 ∼p(x,x + ) 1 B B i=1 ( j̸ =i g(xi,x + j ))δg(xi,x + i )-g(xi,x + i )δ( j̸ =i g(xi,x + j )) g(xi,x + i )( j̸ =i g(xi,x + j ))B→∞ E {(xi,x + i )} B i=1 ∼p(x,x + ) 1 B B i=1 δg(xi,x + i ) g(xi,x + i ) - δ( j̸ =i g(xi,x + j )) j̸ =i g(xi,x + j ) = lim B→∞ E {(xi,x + i )} B i=1 ∼p(xi,x + ) 1 B B i=1 δg(xi,x + ) g(xi,x + ) p(x + | x i ) - k̸ =i δg(xi,x + ) g(xi,x + )-g(xi,x + k )+ j̸ =i g(xi,x + j ) p(x + ) dx + = lim B→∞ E {(xi,x + i )} B i=1 ∼p(x,x + ) 1 B B i=1 δg(x i , x + ) p(x + |xi) g(xi,x + ) -E {(xi,x + i )} B i=1 1 j̸ =i g(xi,x + j ) as B → ∞ p(x + ) dx + = lim B→∞ E {(xi,x + i )} B i=1 ∼p(x,x + ) 1 B B i=1 δg(x i , x + ) p(x + |xi) g(xi,x + ) -E {(xi,x + i )} B i=1 1 j̸ =i g(xi,x + j ) ≜ k(xi) indep. of x + p(x + ) dx + = lim B→∞ E {(xi,x + i )} B i=1 ∼p(x,x + ) 1 B B i=1 δg(x i , x + ) p(x + |xi) g(xi,x + ) -k(x i )p(x + ) dx + = δg(x, x + ) p(x + |x) g(x,x + ) -k(x)p(x + ) dx + .This is zero wheng(x, x + ) = p(x | x + ) k(x)p(x) , i.e., f (x, x + ) = log p(x | x + ) -log c 1 (x) k(x)p(x).as in Eq. ( 10), and Eq. ( 11) follows similarly, exchanging x and x + .Fig. 7 visualizes the inferred waypoints from the task in Fig. 5 .Fig. 8 Justification: This paper provides theoretical analysis of representations learned by contrastive learning. This analysis suggests new ways of using these representations, for both beneficial and harmful purposes. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: N/A Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: See citations to the D4RL dataset.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.Code: https://github.com/vivekmyers/contrastive_planning * Equal contribution.Following prior work[6,10], we will use planning to refer to the problem of inferring intermediate states, not to refer to an optimal control problem.This probabilistic model is equivalent to a discretized Ornstein-Uhlenbeck process[88] and is also known as an AR(1) model[11, Eq. 3.1.16].This can be extended to controlled Markov processes appending the previous action to the observations.While the result of Ma and Collins[54] has C(x) depending on x, the symmetrized version[71] removes the dependence on x.∥ϕ(x0)-ψ(x)∥ 2 2 = p(x | x 0 ) p(x + )C(5) 
 
 Appendix:  


Decision: While the paper's scores alone are borderline, the consensus I gather from reading reviews is that the paper should be accepted, with reviewers praising the paper's clear writing and various technical strengths, particularly after various clarifications given on rebuttal.

Reviews: review: paper_summary: paper_summary: This paper addresses the challenges of probabilistic inference with high-dimensional time series data. It provides compact, closed form solutions in terms of learned representations by extending a prior work to show that the learned representations can be modeled as a Gauss-Markov chain. It demonstrates the effectiveness of this theoretical result via numerical experiments on 39-dimensional and 46-dimensional tasks.
strength_weakness: strength_weakness: The writing of this paper can be improved, particularly in the preliminary section, which is hard to follow due to a lack of self-contained explanations for some notations and definitions. For example, why use product of the marginal distributions to sample negative pairs and what does ≈ mean for similarity? What is the infoNCE objective without resubstitution, particularly what is f and B in (2)? Why the subscript of expectation in (2) seems to not match with the sampling rule (x, x+) ∼ p(x, x) mentioned earlier? What does TSNE in the experiment section stand for?
questions: questions: The usefulness of the geometric properties remains somewhat abstract. While the special case provides some insights, it is not fully clear what it actually means for practical applications. Could you elaborate on how the geometric properties are used in practical scenarios and whether it has been used in your numerical experiments? Why linear interpolation can be used for your high-dimensional tasks while nonlinear interpolation is required for your synthetic task?

I suggest separating Assumption 2 into the assumption inherited from the prior work and a statement of its results. Additionally, the term "prior work" is used throughout the paper, creating some ambiguity. It would enhance clarity by specifying which particular prior work is the one that shows "the representations learned by contrastive learning encode a probability ratio".

Addressing the above two concerns can help understand the contributions better.
limitations: limitations: Yes, the authors have adequately addressed the limitations.
paper_summary: paper_summary: This paper proposes to use contrastive learning technique in time series feature learning, which are helpful for improving inference performance later on.
strength_weakness: strength_weakness: Seeing the analysis of the learned feature unrolling as a Gauss-Markov process reminds me of the very classic Kalman filter. In fact, although Kalman filter is not designed to be learned in the discriminative way that this paper proposes, the equations in this paper share fair bit of similarity to those in Kalman filter. I would like to see some discussion on the difference that separates from this work from the prior time-series modeling with similar formulations.
questions: questions: Does such a contrastive learning method for time series require sampling from the environments repeatedly from any given starting state? If so, how is the applicability of this method on the environments that are only sampled once?
limitations: limitations: N/A
paper_summary: paper_summary: This paper shows that learning contrastive representations with the infoNCE objective combined with an L2 constraint on these representations permit Gaussian distributions on the representations. The authors then show for timeseries data, that, when we make a specific Markov chain assumption (a kind of linear state-space model) then the distribution over learned representations permits a nice Gaussian form for all future or past representations. As a consequence of this, planning can be performed through simple linear interpolation within the learned space.

The paper shortly validates their method in a didactic experimental setup.
strength_weakness: strength_weakness: ### Summary

I find the discussion and narrative of this paper too weak to confirm the novelty and contributions of the method. The authors make assumptions on their model such that it satisfies a nice practical property: planning through interpolation. However, the theoretical results that are presented should be well known consequences of Gaussian assumptions for linear state space models. The presentation makes this insight almost seem profound, but the weak discussion of related work makes me skeptical of the contribution. 

Although this paper presents a nice observation, taking all listed points of improvement into account, I don't think this paper is sufficient yet for an accept. But I'm open to improve the score if the authors address my concerns.

### Major comments
 - Lemma 2 is stated as the main result. Although this way of defining waypoint conditionals is an interesting perspective on planning, the result should be a well-known consequence of using Gaussians (The same goes for Lemma 3). This makes it unclear what the novelty/ contribution of this result is. The authors need to improve this by discussing prior work/ results after introducing their model.
 - The title of the paper: inference via interpolation, is too hidden within the main text (at the end of 4.4 in Eq. 7). I would strongly suggest turning this result into a corollary (indeed with proof referring to Newman and Todd 1958 Pg. 471), then moving this result all the way to the top of section 4, and restructuring the text to move towards this result. 
 - The interpolation in Eq. 7) only presents the mean-representation. There is no such discussion on the covariance.
 - The experiment of figure 5 is not well-designed. The introduced method makes use of strong assumptions that allow it to make use of linear-interpolation for planning towards future or past states. Still, the authors compare this to baselines where (also without discussion) we suddenly assume that these have the same functionality, why is this a reasonable assumption? As a result, a baseline without planning outperforms the baselines **with** planning.
 - Discussion Line 347 (last sentence) strongly oversells the presentation of the paper. The interpolation property is a nice consequence of the method's assumptions (Eq 3. and As.1, 2) and the contrastive learning. However, there is no discussion on how and why other methods cannot achieve similar things. My suggestion is to drop the 'when' comment, and only keep the 'how', because that's what the paper ultimately shows.

### Technical comments
 - Eq 2. is confusing since it is first written that $(x, x^+) \sim p(x, x)$ but in the Eq. it appears as $p(x)$. The term $p(x, x)$ is wrong, write this as $p(x, x^+)$ and keep its use consistent. Perhaps use $(x^+, x) \sim p(x^+ | x)p(x)$ and $x^- \sim p(x)$?
  - Assumption 2, phrasing this result as an assumption weakens the story of the authors, especially since this is the equation that the authors' rely on for later analysis. It's usual to make assumptions or simplifications of the real world, so the fact that approximation error or model misspecification might not satisfy the result should not take away from it. I suggest making this a proposition and writing an assumption around this that discusses approximation error, please directly refer to the proof afterwards (citation + where to find it).
  - Lemma 3, $\eta$ isn't clearly shown to have zero entries apart from the first and last index. Why not use the same definition as in Appendix A.4?
  - Missing proof of Lemma 3, no reference to appendix.

### Minor comments
 - Line 166 typo: 'that that'
- The two representations are $\phi$ and $\psi$ composed with $A$. Calling this two separate representations is confusing, there is only 1 learned representation and $A$ is used to learn the transport from embedding $\psi(a)$ to $\psi(b)$. Maybe state this upfront, i.e., in the first sentence line 187 say $\psi(x)$ and $\phi(x) = A\psi(x)$, that helps when reading chronologically.
 - Line 215, refer back to Eq.3 for the definition of $c$
 - Line 231, proof of lemma 2 is in appendix A.3, not A.2.
 - The three examples in 4.3 for lemma 2 are nice to have for an unfamiliar audience, but I don't see how this contributes to the result since $A$ is learned, and will likely never satisfy trivial solutions.
 - Line 268, the interpolation result is nice, but the linear interpolation insight is not particularly profound (the text makes it seem as if it were so).
 - The start and end symbols in Figure 4 are too small to see and overlap with the black dots too much, even at 500% zoom.
 - Figures 5, 6 are not chronologically ordered/ numbered (6 appears before 5).
 - The choice of baselines and experiments are somewhat limited. I would have liked a comparison to a decoding method which the authors criticize in the introduction paragraph 2 (e.g., using the DreamerV2/ PlaNet RSSM model).
 - There is still space left on page 9, the authors could have spent more space on discussing the limitations of their method, which right now is rather shallow.
questions: questions: .
limitations: limitations: 1) the paper concludes that it provides a kind of 'recipe' for how to get models that have a linear interpolation property for planning, but aside from the model that the authors designed in a very specific way this is not discussed.

2) the experiments are somewhat trivial since the authors expected baselines, that were not designed for a particular feature (planning through linear interp.) to suddenly have this feature.
paper_summary: paper_summary: This paper presents a variant of contrastive learning to time series data. Contrastive learning has been used in many areas but the insight from this paper is joint distribution of representations is also Gaussian. In addition, the experimental results show that their theory can be applied to tasks up to 46 dimensions.
strength_weakness: strength_weakness: 1. Using contrastive learning to address various problems is not novel. However, this paper provides some new insights such as representations learned via temporal contrastive learning following a Gauss-Markov chain.
2. The evaluation section is too simple, and it is possible to add more evaluation. Since it is an ML conference, more experimental results can make the paper more solidated.
questions: questions: 1. As for section 5 Numerical Simulation, is it possible for the authors to use real datasets? For example, the authors can collect some datasets from the Traffic and Stock areas and evaluate the solution mentioned in this paper over them.
2. In addition, I think efficiency is also important for high-dimensional tasks. Could the authors show the running time of each experiment?

Minor:
1. As for citation, the template of NeurIPS is [10, 12] instead of (10;12). I am not sure whether you have changed the template, but it would be great if you could fix it.
limitations: limitations: I think the authors have adequately addressed the limitations.

Comments: title: A solid paper given the finetuned presentation
comment: Thanks for the authors' clarifications and I have no more concerns. This is a solid paper overall, though there is room to enhance the clarity of the writing. Based on these considerations, I have increased my score.
Dear Reviewer,

> This is a good paper with an interesting message that I think should be explored more

Thanks for the kind words :)

> I think a minor rephrasing and 1-2 additional sentences that ...
> add an explicit part to the paper that discusses the rotation matrix ...

We will make these changes.
I want to thank the authors for addressing most of my comments and the additional experiment in the global response.

---
I'm happy that the authors aim to improve upon so many points, all of my concerns are answered. To respond to a few answers in particular:

> Would it be more clear if we separated Lemma 2 into two results?

The main problem I had was that I couldn't exactly deduce what the contribution was, so this concern is about presentation. The sentence leading up to Lemma 2 also read: "Our main result is that the posterior distribution over waypoint representations has a closed form solution ...", which added to this confusion.

I think a minor rephrasing and 1-2 additional sentences that: a) discuss that the result in lemma 2 is a known consequence of the Gaussians, b) cite prior work, and c) that your contribution is that you tie these puzzle pieces together so that you can use this for intermediate state-planning, would clear this concern. This also goes for lemma 3.

> When the matrix $A$ is learned on the datasets used in the paper, we empirically observe that it is often close to a rotation matrix

This is an interesting insight of the method, could the authors add an explicit part to the paper (either main or appendix) that discusses this?

---
I will raise my score from 4 to 7 and contribution from 2 to 3. This is a good paper with an interesting message that I think should be explored more in reinforcement learning literature. I think these types of model-features could potentially improve the scalability of planning algorithms when used in combination with neural networks (e.g., as in EfficientZero).