Contrastive Representations for Temporal Reasoning
Authors: Alicja Ziarko
Michał Bortkiewicz
Michał Zawalski
Benjamin Eysenbach
Piotr Miłoś
Venue: NeurIPS 2025 poster ()
Keywords: Contrastive Learning
Representation Learning
Combinatorial Reasoning
Primary area: deep_learning

Abstract: In classical AI, perception relies on learning state-based representations, while planning --- temporal reasoning over action sequences --- 
is typically achieved through search. We study whether such reasoning can instead emerge from representations that capture both perceptual and temporal structure.  We show that standard temporal contrastive learning, despite its popularity, often fails to capture temporal structure due to its reliance on spurious features. To address this, we introduce Contrastive Representations for Temporal Reasoning (CRTR), a method that uses a negative sampling scheme to provably remove these spurious features and facilitate temporal reasoning. CRTR achieves strong results on domains with complex temporal structure, such as Sokoban and Rubik’s Cube. In particular, for the Rubik’s Cube, CRTR learns representations that generalize across all initial states and allow it to solve the puzzle using fewer search steps than BestFS — though with longer solutions. To our knowledge, this is the first method that efficiently solves arbitrary Cube states using only learned representations, without relying on an external search algorithm.

Contrastive Representations for Temporal Reasoning
Alicja Ziarko 23
Michal Bortkiewiez'
Michal Zawalski!;
Benjamin Eysenbach
Piotr Mitoš 3
University of Warsaw
IDEAS NCBR
3IMPAN
Warsaw University of Technology
Princeton University
NNIIIA
aa ma.ziarkouuv.. .dw..ddupp.
Abstract
In classical AI, perception relies on learning state-based representations, while planning temporal reasoning over action sequences - is typically achieved through search. We study whether such reasoning can instead emerge from representations that capture both perceptual and temporal structure. We show that standard temporal contrastive learning, despite its popularity, often fails to capture temporal structure due to its reliance spurious features. To address this, we introduce Contrastive Representations for Temporal Reasoning (CRTR), a method that uses negative sampling scheme provably remove these spurious features and facilitate temporal reasoning. CRTR achieves strong results on domains with complex temporal structure, such as Sokoban and Rubik's 's Cube. In particular, for the Rubik's 's Cube, CRTR learns representations that generalize across all initial states and allow it to solve the puzzle using fewer search steps than BestF$ - though with longer solutions. To our knowledge, this is the first method that efficiently solves arbitrary Cube states using only learned representations, without relying on an external search algorithm.
Website: https://ppincceoon-tt/ tithhthuu.........o..o.... CR/CRTR/
Introduction
Machine learning has achieved remarkable progress in vision [52], control [21], and language modeling [66, 33]. However, it still falls short on tasks that require structured, combinatorial reasoning. Even relatively simple problems, such planning in puzzles or verifying symbolic constraints, remain challenging for end-to-end learning systems [61, 41]. The best methods for solving these problems use computationally expensive search algorithms. such as A* Best First Search (BestFS) [29].
This work centers on the question: Can learn representations that reduce eliminate the need for search in combinatorial reasoning tasks? To see how good representations can reduce test-time search, consider checking whether graph has an Euler path. This is true if and only if the graph has exactly zero or two vertices of odd degree. A representation that encodes vertex degrees reduces the task to a simple check-elinee theck-eliminating the need for graph traversal. We approach our question by leveraging temporal contrastive learning [46, 56, 21, 19, 44]. These self-supervised techniques designed acquire compact, structured representations that capture the problem's 's temporal dynamics, enabling efficient planning directly within the latent space.
While (Contrastive Learning) CL has shown promise in control tasks [59, 6]. particularly through methods like Contrastive Reinforcement Learning (CRL)[21], we observe that its effectiveness in
Now at Google.
'Equal advising contribution.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

CRL
CRTR
Context 2
Context
Figure 1: CRTR makes representations reflect the temporal structure of combinatorial tasks. 1-SNE visualization of representations learned by CRTR (right) and CRL (left) for Sokoban. Colors correspond to trajectories; three frames from two trajectories are shown the center and linked to their representations. CRL embeddings cluster tightly within trajectories, failing to capture global structure and limiting their usefulness for planning. In contrast, CRTR organizes representations meaningfully across trajectories and time (vertical axis).
combinatorial domains is significantly limited. Specifically, we identify a critical failure mode where contrastive representations overfit to instance-specifie context, instead of reflecting environment dynamics. Consequently, models fail to adequately capture the temporal structure that is vital for effective decision-makin Sokoban (Section 4.1) - manifests as a collapse trajectory representations into small, disconnected clusters, illustrated in Figure (left).
This fallure mode such as when the model overfits to wall layouts in
To solve this, we introduce Contrastive Representations for Temporal Reasoning (CRTR), a simple, theoretically grounded CL method that uses in-trajectory negatives. By design, CRTR forces the model to distinguish between temporally distant states within the same episode. This mechanism prevents it from exploiting irrelevant context such visual layout cues and instead encour- ages learning temporally meaningful embeddings that reflect the problem's 's relevant dynamics. This echoes recent findings in neuroscience, where hippocampal representations of overlapping routes diverge during learning despite visual similarity [10]. Our approach similarly prioritizes temporally meaningful structure over reliance on visual cues.
We evaluate CRTR across challenging combinatorial domains: Sokoban, Rubik's Cube, N-Puzzle, Lights Out, and Digit Jumper. Due to their large, discrete state spaces, sparse rewards, and high instance variability, they serve as challenging testbeds for evaluating whether learned representations support efficient, long-horizon combinatorial reasoning. In each case, CRTR significantly improves planning efficiency over standard contrastive learning, and approaches or surpasses the performance of strong supervised baselines.
Our main contributions are the following:
1. We identify and analyze a critical failure mode standard contrastive learning, showing its inabil- ity to capture relevant temporal or causal structure in problems with complex temporal structure.
2. We propose Contrastive Representations for Temporal Reasoning (CRTR), novel and thcoretically grounded contrastive learning algorithm that utilizes in-trajectory negative sampling to learn high-quality representations for complex temporal reasoning.
3. We demonstrate that CRTR outperforms existing methods out of 5 combinatorial reasoning tasks, and that its representations, even without explicit search, enable solving the Rubik's 's Cube from arbitrary initial states using fewer search steps than BestFS -abbeit with longer solutions.
2 Related Work
We build upon recent advances in self-supervised RL and contrastive representation learning.

---

Contrastive learning. Contrastive learning is widely used for discovering rich representations
from unlabeled data [12, 11, 32, 54] to improve learning downstream tasks [64]. Contrastive learning has enabled effective learning of large-scale models in computer vision [70, 9, 52, 40] and language [60,, 24]. Contrastive learning acquires representations by pulling representations of similar data points, i.e., ones that belong to the same underlying concept, closer together and pushing dissimilar ones further apart in the representation space [63]. This idea is reflected in various contrastive objectives, including Triplet Loss [30], NCE [26], InfoNCE [58], and conditional InfoNCE (CCE) [42]. Our work is most closely related to prior work that uses contrastive learning to obtain representations of time series data [46, 56, 3].
Contrastive representations for sequential problems. Self-supervised contrastive learning has also been applied to sequential (or temporal) problems, including goal-conditioned RL [21, 62, 44], skill-Jearning algorithms [48, 73, 18], and exploration methods [25]. Contrastive representations
also excel at symbolic reasoning for simple mathematical problems [50]. Most temporal-based contrastive algorithms are based on optimizing the InfoNCE objective [58] to distinguish real future states in the trajectory from random states in other trajectories. Interestingly, Eysenbach et al. [19] demonstrate that intermediate state representations can be effectively approximated via linear interpolation between the initial and final embeddings. Based on these findings, we hypothesize that such representations might also reduce or eliminate the need for search in complex combinatorial problems.
Combinatorial problems. Combinatorial environments are characterized by discrete, compact observations that represent exponentially large configuration spaces, often associated with NP- complete problems [35]. Recent RL advancements address these challenges using neural networks to learn efficient strategies, including policy-based heuristics [43, 5], graph neural networks for structural exploitation [8, 36), and imitation learning with expert demonstrations [57]. Our work shows that effective search-facilitating representations can be learned from suboptimal data, without relying on expert demonstrations.
Planning in latent space. Efficient planning complex environments can be achieved by learning state representations that reflect the underlying structure of the problem. Prior work on world models [27, 28] shows that compact representations of high-dimensional environments are critical for agent performance. Another line of research [55, 22] focuses on learning representations that retain only the features relevant for planning. In robotic settings, similar approaches train latent representations to support movement and decision-making [31, 22]. We build on the approach of Eysenbach et al. [21], which formulates coul-conditioned reinforcement learning as a representation learning problem, with the key distinction that we focus on combinatorial reasoning tasks.
3 Preliminaries
Combinatorial problems and dataset properties. We focus on combinatorial problems, which can be formulated as deterministic goal-conditioned controlled Markov processes (S, A. p, Ta 7). At each timestep the agent observes both the state 8 E S and the goal g E S, and performs an action E A. We assume that the transition function A x S - S is known and deterministic. Initial states are sampled from the distribution (o()s). We define reward function (st) 1 for st g and ry (st) = otherwise.
The objective is to learn goal-conditioned policy <a s, g) that maximizes the expected reward: max, Pp.( 1>/= 7'ry at)] We study an offline learning setup with a dataset of success- ful yet suboptimal trajectories т; = (($1, 01), (82, az).... . (g. -)). A state 8, is defined as reachable from 8) if there exist a path a a2...... . Ons such that s, p(a,, p(a, 1. p(.... . p(a,.,,((1,, s1))))
Contrastive reinforcement learning. We employ contrastive reinforcement learning (CRL) method [21] to train a critic f(s g). which estimates the similarity between the current state s and future state g (goal). The critic uses a single encoder network vé to generate both representations of the state (s) and the goal (g)). Critic's 's output measures a similarity between these representations
with a metric f (s, = [[((s) (ggg) that reflects the closeness of the states. For details, see Appendix C. To train the critic, we construct batch B by sampling random trajectories from the dataset. For each trajectory, we select state 8; uniformly and draw a goal gt using a Geom -
distribution over future states. Negative pairs consist of state-goal pairs (st, 9z) with goals sampled
3

---

from different trajectories. We use the critic's outputs in the InfoNCE objective [58], following prior work in CRL [21, 20, 72, 71, 44, 6]. This objective provides lower bound on the mutual information between state and goal representations:
|8|
min EB
log
2j=1 efo(s1.9))
Mutual information. The mutual information between random variables X and Y is the amount of information the value of one conveys about the value of the other. Formally, this corresponds to how much the entropy decreases after observing of the random variables: T(X; Y) H[X] [[[ Y]. The conditional mutual information T(x:Y C) measures the extra information Y provides about X once the context C' is known: I(X; C) H[X C] H[X Y, C.C. The conditional MI is zero when X and Y are conditionally independent given C.
Search-based planning. Several prior works use explicit search algorithms solving complex environments [57, 7, 67, 47]. In our study. we focus on the Best-First Search (BestFS) [49]. BestFS builds the search tree by greedily expanding nodes with the highest heuristic estimates, hence targeting paths that are most likely to lead to the goal. In our approach, BestFS uses the representations of the current state and goal, and relies on the critic to evaluate the distances between neighboring states and the goal. While not ensuring optimality, BestFS provides simple yet effective strategy for navigating complex search spaces. The pseudocode for BestFS outlined in Appendix B, In our work, we use distances in the latent space as the heuristic, as detailed in Section 3.
4 Learning Temporal Representations that Ignore Context
The main contribution of this paper is a method for learning representations that facilitate planning. We start by describing how naive temporal contrastive representations fail in combinatorial problems. Using Sokoban as an example, we will highlight why this happens (Sec. 4.1), and use it to motivate (Sec. 4.1) a different contrastive objective that better facilitates planning on many problems of interest. Section 4.3 summarizes our full method, CRTR.
Problem definition. We use a neural network : S R* to embed states into k-dimensional representations. We define the critic f((. g) [|0(s) ا(9)) as the norm between these learned representations (see Appendix C). Our aim will be to learn effective representations from dataset of trajectories ($t).=L..N, collected from either suboptimal expert or random policies. We evaluate these representations by testing whether their distances reflect the structure needed to solve combinatorial tasks like Rubik's Cube or Sokoban.
4.1 Failure of Naive Temporal Contrastive Learning in Combinatorial Domains
A straightforward approach to learning representations o(s) is to use temporal contrastive learning [46, 56] (outlined in Sec. 3): positive pairs are sampled from nearby states within the same trajectory, and negatives from different trajectories. However, applying this approach to a common benchmark (Sokoban) reveals an important failure mode. This section introduces this failure mode and its mathematical underpinnings, Sec. 4.2 introduces an idealized algorithm for fixing this failure mode, and Sec. 4.3 turns this idealized algorithm into a practical one that we use for our experiments in Sec. 5.
Sokoban is puzzle game where agent must push boxes target locations in a maze. Each level (or problem instance) is generated with a random wall pattern that is different from one trajectory to another (see Fig. 1). Fig. I shows t-SNE projection of representations learned by temporal contrastive learning on this task. We observe that embeddings from different mazes form tight, isolated clusters. Appendix Fig. 15 shows the same phenomenon happens on the Digit Jumper task. Thus, the representations primarily encode the layout of the walls and not the temporal structure of the task. The reason representations use those features is that they minimize the contrastive objective. Each batch element typically comes from a different maze, so representations that use the wall pattern to detect positive vs negative pairs achieve nearly perfect accuracy. Thus, we will need a different objective to learn representations that primarily focus on temporal structure, and ignore static features (like the walls Sokoban). To this, we will first provide a mathematical explanation for this failure mode, which will motivate the new objective in Sec. 4.2.

---

A mathematical explanation. The failure of temporal contrastive learning can be explained with context variable c:
Definition 4.1. Let D = {TI, T2...., TnJ be a dataset of trajectories, where each trajectory T ($1...., ST) ) is sequence of states. Assume that each state can be written s; (c,, fi), where oa is a component that remains constant over the whole trajectory: C1 = ... = Then we can write the trajectory (c. (fi..... fr)). We refer to as the context variable, and (fa...., fr) as the temporal part of the trajectory.
For instance, in Sokoban, in each trajectory. box goals and the temporal part to be the positions of the player and boxes. It is always possible to decompose trajectory into context and temporal part by setting c = (). but we will be primarily interested in cases where the context is disjoint from the temporal part:
take the context to be the positions of walls and
Assumption 4.2. Let be context and ($1..... .8) ~ D trajectory. Then for any time steps i < J, the state 8y and the context e are conditionally independent given 8y: 1 c 8;-
This Assumption holds for Sokoban, where the context is the wall layout, since the wall layout can be fully determined by each of the states in the trajectory. We will use this assumption motivate an idealized algorithm in Sec. 4.2, but our practical method in Sec. 4.3 will not require this assumption. Indeed, our experiments (Section 5) show that the approach also works when the context evolves slowly over time (e.g., for the Rubik's Cube).
4.2 Learning Representations that Ignore Context: An Idealized Algorithm
Based on this mathematical understanding of the context, we now introduce an idealized alternative method for learning the representations, which makes use of information about the contexts. Sec. 4.3 introduces the practical version of this algorithm, which does not require any priori information about the context.
The key idea in our method is to sample negative pairs (x, x-) ) that have the same context, so that the context features are not useful for distinguishing positive and negative pairs; thus, these context features will not be included in the learned repersentations. Specifically, our idealized method works by first sampling the context c P(C"), then positive pairs (xx, 2+) from the conditional joint
distribution P(X, X+ c) and negatives from the marginal conditional distribution x ~ P(X |c)
for E {I....,N 1}. The resulting contrastive learning objective is:
(
max C(f) A 1, -(x(e
N
( EN e ((
(
max L(f) 2 f
E (ag,ay) J-
N
(( + N- el
P(X)
where .(..) is the learned similarity score between state pairs. For example, in Sokoban, this method changes how negative examples sampled so that they always have the same wall configuration.
Mathematically, this alternative objective a lower bound on the conditional mutual information T(X; X4 C') [42], which can be decomposed using the chain rule for mutual information:
T(X+; X C) = I(X+; C X) + 1(X+; X) - I(X+; C).
Combined with the assumption that X+ 1 C (so / 1(x,;C X) = 0), we see that this new objective is lower bound difference of mutual informations:
L(f) < I(X: C()+;C).
The term I(X; X4) is what is usually maximized by temporal contrastive learning. The second term, I(X+: C), is akin to adversarial feature learning [38, 16, 65], which aims to learn representations that do not retain certain pieces of information. Our method achieves similar effect without the need for adversarial optimization. This context-invariance objective parallels mechanisms observed in biological systems, where overlapping episodie memories are actively decorrelated to reduce interference [10].
5

---

Algorithm 1 CRTR performs temporal contrastive learning, but samples negatives in a different way so that representations discard fask-irrelevant context, boosting performance (See Fig. 2).
dataset. shape [num_ traj, traj_ len, dom_
t0 = up. random. (dataset. Sadataase.ssappeei [1], batch_size)
t1 = t0 mp. random. geometric(1 - discount, batch_size)
traj_id _ii np- random. www.om.choice choice (dataset shapee batch_size)
new line of code for CRTR (our approach);
traj_id mp ap.repeat (traj_id[: repetition_facttnt axis=0)
// repetition_factor]]
batch (dataset (traj_id, t0], dataset(traj_ic t1])
further batch processing, the same for CRL and CRTR
4.3 A Practical Method
While the idealized method in Sec. 4.2 is useful for analysis, the main challenge with practically implementing this method is that the context is not clearly separable from the observation. For example, upon seeing a Sokoban board for the first time, how should one know that blocks are movable (not part of the context) while walls are not (should be part of the context)? This section proposes a practical algorithm that does not require the assumption from Definition 4.1; our experiments in Sec. 5 will demonstrate that this practical method continues to work when this assumption is violated.
The key idea behind our practical method changes how training pairs are sampled for contrastive learning: for each trajectory included in batch, sample multiple positive pairs. Contrastive learning treats each batch element as forming negative pairs with all other elements in the batch. When multiple positives from the same trajectory are present, some negative pairs will consist of two states from the same trajectory- likely sampled from different points in time. These within-trajectory negatives are drawn from a different distribution than the corresponding positives, often involving greater temporal separation. As a result, the model is encouraged to focus on temporal distinctions rather than features that constant throughout trajectory.
Implementing this idea in practice requires changing just a few lines of code from prior temporal contrastive learning methods, as highlighted in Algorithm 1). The repetition factor governs the proportion of such negatives, thereby providing controllable mechanism interpolate between the standard and proposed objectives. Using data sampled in this way guarantees that some negative training pairs in each batch come from the same trajectory. We compare with potential alternative approaches in Appendix I.
4.4
What if the context is not constant?
While our theoretical analysis required that the context be clearly separable from the observation, the key insight behind our practical method (Sec. 4.3) was to lift this assumption, allowing the method to be applied without any knowledge of the context, even to problems without a constant context (e.g., the Rubik's Cube).
While we empirically test how our method performs on problems without fixed context in See. 5, here provide some intuition for why the method might be expected to work in such settings. In the Rubik's 's Cube, for instance, all states are mutually reachable, making it difficult to define the context. Nevertheless, if we focus the more shuffled portion of the trajectory, simple heuristics like Hamming distance can classify state pairs with 90% accuracy. This suggests that while each move introduces temporal change in some parts of the cube, others remain unchanged, implicitly forming type of context. As a result, networks may latch onto features that correlate with this pseudo-conted rather than true temporal proximity. In Section 5.2 we empirically demonstrate that method can also improve performance in settings without a constant context.
5 Experiments
Our experiments aim to answer the following specific research questions:

---

1. Does learning representations that ignore context improve performance on combinatorial reason- ing problems? (Sec. 5.2)
2. Do learned representations alone suffice for reasoning, or is explicit search essential? (Sec. 5.3)
3. Are representation learning methods that remove context competitive with successful prior methods for combinatorial reasoning? (Sec. 5.2)
4. What is the relative importance of design decisions, such as how the negatives are sampled and the number of in-trajectory negatives? (Sec. 5.4)
5.1 Experimental Setup
Environments. We evaluate all methods on five challenging combinatorial reasoning tasks: Sokoban [17], Rubik's Cube, N-Puzzle [34], Lights Out [2], and Digit Jumper [4]. Most of these are NP-hard [15, 13, 53] and serve as standard RL benchmarks [1, 51, 69]. Sokoban is a grid-based puzzle where an agent pushes boxes targets while avoiding irreversible states. Rubik's Cube requires aligning each face of a 3D cube to single color. N-Puzzle involves sliding tiles within 4 x 4 grid to reach a goal configuration. Lights Out is toggle-based puzzle aiming switch all cells to an off state. Digit Jumper is a grid game where each cell indicates the jump length from that position. See Appendix for full environment details.
Baselines. We compare against three baselines. The contrastive buseline performs temporal
contrastive learning [56, 46, 21], training representations without in-trajectory negatives. We will refer to this baseline as contrastive RI. (CRL. [21]. The supervised baseline [14, 68] predicts state distances using value network trained via imitation on demonstrations. Finally, the DeepCubeA [i] baseline learns a value function via iterative one-step lookahead. As lower bound, we also evaluate the performance of representations from a randomly initialized network. For fairness, the CRL and supervised baselines use the same architecture as CRTR.
We will evaluate methods in two settings: with and without search. When we use search, all methods, including DeepCubeA, use BestFS for planning. During tree search, all actions are considered for Rubik's ' Cube, N-Puzzle, Digit Jumper, and Sokoban; for Lights Out, expansion is limited to the top ten actions ranked by the value function. All the methods avoid loops by only considering states that were not already processed. Further evaluation details are provided in Appendix D. When evaluating methods without search, we greedily find the neighboring states with minimum predicted distance and select the action leading to that state (recall, we assume dynamics are known and deterministic).
Metrics. We evaluate two key aspects of representation quality. Spearman's 's rank correlation measures the aligmment between representation-te space distance and the actual number of steps each state is from the last state in trajectory. We compute this for each trajectory in the test set, then average the results over 100 trajectories. A high correlation indicates that states closer in time are also closer in the representation space. We demonstrate that this correlation is good measure of representation quality in Appendix G. Success rate at a fixed computational budget measures the fraction of initial states from which complete solution is found. This metric reflects whether the learned representations support effective planning.
The training details, including hyperparameters, network architectures, and dataset descriptions, specified in Appendix C. Code to reproduce our experiments is available online: Com/Princeton RRL/CRLERRRRRE
5.2 Context-Free Representations for Combinatorial Reasoning
We first analyze the representations learned by CRTR visually, using Sokoban as a testbed because it contains obvious context features (the wall positions). Fig. shows the representations learned by CRTR and compares them with those learned by standard temporal contrastive learning (CRL), which differs from our method by not using in-trajectory negatives. We use t-SNE to reduce the representations down to 2D. The CRL baseline learns representations that cluster together, likely focusing on static features for each trajectory, all observations get encoded to very similar representations. In contrast, the representations from CRTR likely focus on temporal structure - the fact that observations from different trajectories but similar stages of solving get encoded similar representations indicates that the representations are ignoring context information (which is irrelevant for decision making).
7

---

Rubik
Lights Out
Sokoban
1.2
1.0
Success Rate
0.5
0.0
0.0
0.0
5000 Search Budget (Nodes)
2500
0
5000 Search Budget (Nodes)
2500
2500
5000
Search Budget (Nodes)
Jumper
15-puzzle
Success Rate
0.8
1.0
CRTR
CRL
0.5
Supervised
0.4
DeepCubeA
Random
200
400
1000
2000
Search Budget (Nodes)
Search Budget (Nodes)
Figure CRTR performs well in evaluated domains. Success rate as a function of search budget across five domains. CRTR compared to baselines: CRL [21], Supervised [14] and DeepCubeA [1]. Results are averaged over seeds; shaded regions indicate standard error. For the Rubik's 's Cube, both the supervised and random baselines achieve success rate of zero for all search budgets.
Our second experiment studies whether these CRTR represen- tations are useful for decision making. do this, we use representations to construct heuristic for search, measuring the fraction of problem configurations that are solved within given search budget (X axis). As shown in Figure 2, CRTR consistently achieves among the highest success rates, with the largest gains in Sokoban and Digit Jumper. The strong performance relative to CRL highlights the importance of re- moving context information from learned representations. In Appendix E, we provide additional, smaller- scale experiments showing that these improvements also hold when using non- greedy solver.
Sokoban Correlation
Correlation
CRTR
CRL
30
Evaluation Epoch
Figure 3: Distances given by CRTR representations reflect the temporal structure well. Correlation (Spear- man's p) between the distance induced by learned embeddings and actual dis- tance across the training, CRTR com- pared with CRL.
Our third experiment compares CRTR to supervised ap- proaches [1, 14] for solving combinatorial problems. Again, Fig. 2. CRTR ranks among the hop-pefforming methods in each environment and is strictly the best in two environments. In contrast, the supervised baselines perform much worse in Rubik's 's Cube and Lights Out. We conjecture that the improved performance of CRTR come from how it represents values as distance between learned representations, rather than number output by a monolithic neural network.
The t-SNE visualizations (Figure 1) suggests that CRL focuses primarily on the static context, while CRTR focuses on the temporal structure. Below, we present additional empirical evidence supporting this interpretation.
We perform further analysis in Sokoban environments. Without negative pairs, the classification task becomes nearly trivial: the model leverages context cues to achieve close to 100% accuracy (Appendix E). Despite this, the learned representations exhibit low correlation with ground-trute state-space distances (Figure 3), indicating that the model ignores temporal structure and instead relies on static context. In contrast, CRTR prevents reliance on contextual shortcuts, resulting in representations that better capture the underlying geometry of the environment (Figure 3). We provide a similar analysis for Digit Jumper in Appendix F. We also demonstrate that using CRTR leads to improved temporal structure in robotic domains (See Appendix F). In Appendix H we show that

---

Rubik
Lights Out
Sokoban
Digit Jumper
15-Puzzle
1.0
Success Rate
0.0
15
300
6000 10
200
4000 10
30
90
27
81
120
360
Solution Length
Solution Length
Solution Length
Solution Length
Solution Length
CRTR without Search
CRTR with BestFS
Figure 4: CRTR solves most tasks without requiring any search. Adding BestFS results in shorter solutions. We plot the fraction of configurations solved with solution length of at most while limiting the number nodes created to 6000. Surprisingly, on the Rubik's cube CRTR achieves higher success rate without search, solving all board configurations within the budget. Figure 11 in the Appendix E presents the same results, but for CRTR, as well as the Contrastive and Supervised baselines.
CRTR results in representations that optimize conditional mutual information [(X, (IC,), while CRL does not.
5.3 Is search necessary?
Do good representations allow solve combinatorial problems without search, or least reduce the amount of search required to get high success rates? We study this question by using the learned representations to perform greedy planning (see Sec. 5.1) for up to 6000 search steps.
We present the results from this experiment in Figure 4, showing the fraction of problems solved with fewer than certain number of steps. We compare to the variant of CRTR used in Sec. 5.2 that uses the representations to perform search. As an example of how to interpret this plot, the red line the lefumost plot shows that CRTR (without search) solve about 50% of Rubik's s cube configurations in less than 300 moves. We compare to the CRL and Supervised baselines in Appendix E.
On 4 5 tasks (Rubik's Cube, Lights Out, 15-Puzzle, and Digit Jumper), CRTR solves nearly all problem instances. The key takeaway thus: for most problems, CRTR can find solutions without needing any search at all. Perhaps the most interesting result is the Rubik's s cube, where we found that our representations can solve all problem instances in less than 6000 moves. Surprisingly, using search decreases the total fraction of Cube configurations that solved. The main failure mode is Sokoban, where success rates are low, likely due to the presence of dead-ends in the game.
However, avoiding search comes cost: the solutions found without search are typically longer than those found with search. For example, on the Rubik's Cube the average solution length is around 400 moves (see Appendix E) much higher than the optimal of 26 moves and even higher than the 100 moves that beginner cube solving methods require.
Step 1 Fully scrambled.
Step Assembled lower 2x2x2 block.
Step Assembled 2x2x3 block and aligned smaller 1x1x3 blocks.
Step Almost all pieces in correct configuration.
Step Solved!
Figure 5: CRTR without search exhibits bock-building-liae like behavior. Intermediate states from solving a randomly scrambled cube, illustrating how the algorithm gradually builds partial structure. The average solve is about 400 moves, and we see similar block building behavior across solves.

---

This simple greedy approach - just picking the neighbor closest to the goal starts to show hints of algorithmic behavior. On Rubik's 's Cube, for example, it learns something that looks like a rudimentary form of block-buildigg (See Fig. 5, a common strategy used by humans for solving the cube. This block building strategy was not programmed or explicitly rewarded, but instead emerged from training the representations on random data.
5.4 Additional Experiments
Appendix I presents additional ablation experiments. We find that (1) our strategy for sampling data (Alg. 1) outperforms several alternatives, and (2) CRTR is robust to the repetition hyperparameter, with 2 being a good choice in all settings we have tested.
We report several unsuccessful attempts and negative results in Appendices K and E.
Conclusion
One way of reading this paper is story about search. While much of the early history of AI was built upon search, and while search has regained popularity in recent years as an effective strategy for sampling from language models, our paper shows that many sophisticated reasoning problems can be solved without any search. The key catch is that this is only true if representations are learned appropriately (by removing static context features). While we have only demonstrated results on combinatorial problems, which have known and deterministic dynamics, we look forward to extending these techniques to problems such as chemical retrosynthesis and robotic assembly problems that have a rich combinatorial structure, but which introduce additional complexity because of unknown and stochastic dynamics.
Another way of reading this paper is as a story about the value of metric embeddings. Image classifiers automatically identify patterns and structures in images (arrays of pixels), mapping images to representations so that semantically-simile representation learning method does the same for combinatorial reasoning problems, automatically identifying patterns and structures of observations so that the simplest possible decision rule (greedy action selection) can solve some of the most complex reasoning problems (e.g., Sokoban in P-SPACE complete [13]). The success of these representations raises the question of how much reasoning can be done in the representation space. Can we think about compositional reasoning as vector addition? Can we think of length generalization as rays in representation space? And, perhaps most alluringly. can we exploit representational geometry to use the solutions simple problems to find solutions to hard problems, by using local rules to provide global structure on the space of representations?
images are mapped to similar representations. Our
Limitations Our results characterize the behavior of contrastive methods deterministic combina- torial environments. However, this work does not account for stochastic settings. Many real-world applications, such as robotic manipulation and autonomous driving, are inherently stochastic, so our findings may not directly generalize to these domains. Moreover, we focus exclusively on fully observable environments with easily identifiable contextual information the observations. Evaluating whether our approach can be extended to partially observable environments is left for future work.
Acknowledgments. The authors would like to thank Raj Ghugare and Mateusz Olko for helpful discussions and feedback. MB was supported by the National Science Centre, Poland (grant no. 2023/51/D/ST6/016 and the Warsaw University of Technology through the Excellence Initiative: Research University (IDUB) program. MZ was supported by the National Science Centre, Poland (grant no. 2023/49/NSST6/0S 19). We gratefully acknowledge Polish perfoorrffrm computing infrastructure PLGrid (HPC Center: ACK Cyfronet AGH) for providing computer facilities and support within computational grant no. PLG//024/017382 We acknowledge Polish high-performance computing infrastructure Pl.rrid for awarding this project access to the LUMI supercomputer, owned by the EuroHPC Joint Undertaking, hosted by CSC (Finland) and the LUMI consortium through PLL/2024/060011136
10

---

References
[1] Agostinelli, F,, McAleer, S., Shmakov, A., and Baldi, P. (2019). Solving the rubik's cube with deep reinforcement learning and search. Nat. Mach. Intell., 188)3356-363.
[2] Anderson, M. and Feil, T. (1998). Turning lights out with linear algehra. Mathematics magazine, 71(4):300- 303.
[3] Assran, M., Duval, Q. Misra, Bojanowski, P., Vincent, P., Rabbat, M., LeCun, Y., Ballas, N. (2023). supervid learning from images with a jmbeddimm predictive architecture. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 15619-15629
[4] Bagatella, M., Olšák, M., Rolínek, M., and Martius, G. (2021), Planning from pixels in environments with combinatorially hard search spaces. Advances Neural Information Processing Systems, 34:24707-24718.
[S] Bello, I., Pham, H., Le, Q. V. Norouzi, M., and Bengio, (2016). Neural combinatorial optimization with reinforvement learning. CoRR, 1oss1611.099400
[6] Bortkiewicz, M., Palucki, W. Myers, V. Dziarmaga, T., Arczewski, T., Kuciński, L., and Eysenbach, B. (2024). Accelerating goal-conditioned rl algorithms and research. arXiv preprint arXiv:2408. 11052.
[7] Brown, N., Bakhtin, A., Lerer, and Gong, Q. (2020). Combining deep reinforcement learning and search for imperfect-infoomation games. Larochelle, H., Ranzato, M., Hadsell, R.. Balcan, M., and Lin, H.. editors. Advances in Neural Information Processing Systems 33; Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual.
[8] Cappart, Q. Chétclat, D., Khalil, E. B., Lodi, A., Morris, C., and Velickovic, P. (2021). Combinatorial optimization and reasoning with graph neural networks. In Zhou, Z., editor, Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, IJCAI 2021, Virtual Event / Montreal, Canada, 19-27 August 2021, pages 4348-4355 ijcai.org.
[9] Caron, M., Touvron, H.. Misra, I., Jégou, H., Mairal, J., Bojanowski, P., and Joulin, A. (2021). Emerging properties in superviie vision transformers. arXiv preprint arXiv: 2104.14294
[10] Chanales, A. J. H., Oza, Favila, E., and Kuhl, B. (2017). Overlap among spatial memorics triggers repulsion of hippocampal representations. Current Biology, (1115))23772317..0
[11] Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. (2020). A simple framework for contrastive learning of visual representations. In III. H. D. and Singh. A., editors. Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings Machine Learning Research, pages 1597-1607. PMLR.
[12] Chuang, C.-Y,, Robinson, J., Lin, Y-C.,, Torralba, A., and Jegelka, (2020). Debiased contrastive lcarning. Advances in neural information processing systems, 33:8765-877 -777555..
[13] Culberson, J. (1998). Sokoban is pspace-complete. In Proceedingy of the International Conference on Fun with Algorithuss, pages 65-76, Elba, Italy. Carelton Scientific.
[14] Czechowski, K.. Odrzygozdz. T., Zbysinski. M.. Zawalski, M.. Olejnik, K.. Wu, Y., Kucinski, L., and Milos, (2021). Subgoal search for complex reasoning tasks. Ranzato, M.. Beygelzimer. A., Dauphin. Y. N., Liang, P., and Vaughan, W., editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 1,21, December 6-14, 2021, virtual, pages 624-638.
[15] Demaine, E. D., Eisenstat, S., and Rudoy, M. (2018). Solving the Rubik's Cube Optimally NP-complete. In Niedermeier, R. and Vallée, B., editors, 35th Symposium on Theoretical Aspects of Computer Science (STACS 2018), volume 96 of Leibniz International Proceedings in Informatics (LIPIcs), pages 24:1 14:11, Dagstuhl, Germany. Schloss Dagstuhl Leibniz- Zentrum für Informatik.
[16] Donahue, J., Krähenbühl, P., and Darrell, T. (2016). Adversarial feature learning. arXiv preprint arXiv:1605.09787
[17] Dor, D. and Zwick, U. (1999). Sokoban and other motion planning problems. Comput. Geom., 13(4):215- 228.
[18] Eysenbach, B., Gupta, A., Ibarz, J., and Levine, S. (2018). Diversity is all you need: Learning skills without a reward function. International Conference on Learning Representations.
11

---

[19] Eysenbach, B., Myers, V., Salakhutdinov, R.. and Levine, S. (2024). Inference via interpolation: Contrastive representations provably enable planning and inference. Advances in Neural Information Processing Systems, 37:58901-58928.
[20] Eysenbach, B., Salakhutdinov, R., and Levine, S. (2021). C-learning: Learning to achieve goals via recursive classification. In International Conference on Learning Representations.
[21] Eysenbach, B., Zhang, T., Levine, S., and Salakhutdinov, R. R. (2022). Contrastive learning as goal- conditioned reinforcement learning, Advances in Neural Information Processing Systems, 35:35603-35622
[22] Fang, K. Yin, P., Nair, A., Walke, H., Yan, G., and Levine, S. (2022.. Generalization with lossy affordances: Leveraging broad offline data for learning visuomotor tasks. In Liu, K., Kulic, D., and Ichnowski, J., editors, Conference on Robot Learning, CoRI. 2022, 14-18 December 2022, Auckland, New Zealand, volume 205 of Proceedings of Machine Learning Rexearch, pages 106-117. PMLR.
[23] Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, S. (2020). D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprini arXiv:200..07219
[24] Gao, Y., Xiong, Y., Gao, X., Jia, K., Pan, J., Bi, Y., Dai, Y., Sun, J., Wang, M., and Wang. H. (2023). Retrieval-angmented generation for large language models: A survey. arXiv preprint arXiv: 2312.10997
[25] Guo, Z., Thakoor, S., Pistar, M., Pires, B. A., Altché, F., Tallec, C., Saade, A., Calandriello, D., Grill, J., Tang. Y.. Valko. M.. Munos. R., Azar, M. G., and Piot, B. (2022). Byol-explore: Exploration by bootstrapped prediction. In Koyejo. S., Mohamed, S.. Agarwal, A., Belgrave, D.. Cho. K., and Oh. A... editors, Advances Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022.
[26] Gutmann, M. U. and Hyvärinen, A. (2012). contr-conte estimation of unnormalized statistical models, with applications to natural image statistics. Journal of machine learning research, 13(2).
[27] Ha, D. and Schmidhuber. J. (2018). World models. arXiv preprint 1rXiv:1803.1012.
[28] Hafner, D., Pasukonis. J., Ba. J., and Lillicrap. T. (2023). Mastering diverse domains through world models. arXiv preprint 2xXxv:2301.04100
[29] Hart, P. Nilsson, N. J., and Raphael, B. (1968). A formal hasis for the heuristic determination of minimum cost paths. IEEE Trans. Syst, Sei. Cybern, 1(2):100-107.
[30] Hoffer, E. and Ailon, N. (2014). Deep metric learning using triplet network. arXiv preprint arXiv: 1412.6622.
[31] Ichter, B. and Pavone, M. (2019). Robot motion planning in learned latent spaces. IEEE Robotics Autom. Lett. 4434:2007-2414.
[32] Jaiswal, A., Babu, A. R., Zadeh, M. Z., Banerjee, D.,, and Makedon, F. (2020). survey contrastive sclf-supervised learning. Technologies.
[33] Jiang, A. Q., Ziarko, A., Piotrowski, B., Li, W., Jamnik, M., and Milos, P. (2024). Repurposing language models into embedding models: Finding the compute-optimat recipe. In Globersons, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomczak, J. M., and Zhang, C., editors, Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December 10 15, 2024.
[34] Johnson, W. W. and Story, W. E. (1879). Notes on the "15" puzzle. American Journal of Mathematics, 2(4):397-404 404,
[35] Karp. R. M. (1972). Reducibility among combinatorial problems. In Miller, R. E, and Thatcher. J. W.. editors, Proceedings of symposium on the Complexity of Computer Computations, held March 20-22, 1972, at the IBM Thomas J. Watson Research Center, Yorktown Heights, New York, USA, The IBM Research Symposia Series, pages 85-103. Plenum Press, New York.
[36] Kool, W. van Hoof, H., and Welling, M. (2019). Attention, learn solve routing problems! 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenRevicw.net.
[37] Kraskov, A., Stögbauer, H., and Grassberger, P. (2004). Estimating mutual information. Physical review. E, Statistical, nonlinear, and soft matter physics, 69:066138.
12

---

[38] Li, H. Pan, J.J., Wang. S... and Kot, A. C. (2018). Domain generalization with adversarial feature learning. In Proceedings of the IEEE conference on computer vision and patlern recognition, pages 5400-5409.
[39] Li. Z., Guo. and Si, X. (2024). G4satbench: Benchmarking and advancing sat solving with graph neural networks.
[40] Liu, H., C., Wu, Q., and Lee, Y.J. (2024). Visual instruction tuning. Advances in neural information processing systems, 36.
[41] Ma, K., Du, X., Wang, Y., Zhang, H., Wen, Z., Qu, X., Yang, J., Liu, J., Liu, M., Yue, X., ct al. (2024). Kor-bench: Benchmarking language models on knowledge-orthoggnall reasoning tasks, arXiv preprint arXiv:4410.0652..
[42] Ma. M. Q. Tsai, Y.HH. H., Liang, P. P., Zhao, H., Zhang, K., Salakhutdinov, R., and Moreney, L.-P. (2022).
Conditional contrastive learning for improving fairness in self supervipe learning.
[43] Mazyavkina, N., Sviridow, S., Ivanov, S., and Burnaev, E. (2021). Reinforcement learning for combinatorial optimization: A survey. Comput. Oper. Res., 134:105400.
[44] Myers, V., Zheng, C., Dragan, A., Levine, S., and Eysenbach, B. (2024). Learning temporal distances: Contrastive successor features can provide metric structure for decision-making International Conference on Machine Learning.
[45] Nauman, M., Cygan, M., Sferraza, C., Kumar, A., and Abheel, P. (2025). Bigger, regularized, categorical: High-capacity value functions efficient multi-task learners. CoRR, ahs/2505.23 25500
[46] Oord, A. v. d., Li, Y. and Vinyals, O. (2018). Representation learning with contrastive predictive coding. arXiv preprint 1807.03377
[47] Orseau, L., Lelis, L., Lattimore, T., and Weber, T. (2018), Single-agent policy tree search with guarantees. In Bengio, S.,. Wallach, H. M., Larochelle, H., Grauman. K.,, Cesa- Bianchi, N., and Garnett, R., editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montreal, Canada, pages 9205-3215.
[48] Park, S., Rybkin, O., and Levine, S. (2023). Metra: Scalable unsupervised rl with metric-aware abstraction. International Conference on Learning Representations.
[49] Pearl, J. (1984). Heuristics intelligent search strategies for computer problem solving. Addison- Wesley series artificial intelligence. Addison Addison-Wesley.
[50] Poesia, G.. Dong, W., and Goodman, N. (2021). Contrastive reinforcement learning of symbolic reasoning domains. Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang. P., and Vaughan, W., editors, Advances in Neural Information Processing Systems, volume 34, pages 15946- 159666 Curran Associates, Inc.
[51] Racanière, S., Weber, T., Reichert, D. P... Buesing, L... Guez, A.. Rezende, D. J., Badia, A. P., Vinyals, O.. Heess, N., Li, Y., Pascanu, R., Battaglia, W., Hassabis, D., Silver, D., and Wierstra, D. (2017). Imagination- augmented agents for deep reinforcement learning. In Guyon, I,, von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Gamett, R., editors, Advances in Neural Information Processing Systems 30: Annual Conference Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA. USA, pages 5690-5701.
[52] Radford, A., Kim, J. W., Hallacy. C., Ramesh, A.. Goh, G., Agarwal. S., Sastry, G., Askell, A., Mishkin, P.. Clark, J.. Krueger, G., and Sutskever. 1. (2021). Learning transferable visual models from natural language supervision. International Conference on Machine Learning.
[53] Ratner, D. and Warmuth, M. K. (1986). Finding shortest solution the N N extension the 15 15-puzzle is intractable. In Kehler, T., editor, Proceedings of the 5th National Conference on Artificial Intelligence. Philadelphia, PA, USA, August 11-5,, 1986. Volume I: Science, pages 168-172. Morgan Kaufmann.
[54] Rethmeier, N. and Augenstein, I. (2023). A primer on contrastive pretraining in language processing: Methods, lessons learned, and perspectives. ACM Computing Surveys, 15(10)
[55] Schrittweseer, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, I.,, Schmitt, S., Guez, A., Lockhart, E., Hassabis, D.. Graepel, T. Lillicrap. T. P., and Silver, D. (2020). Mastering atari, go, chess and shogi by planning with learned model. Nat., 588(7839):604-699
[56] Sermanet, P., Lynch, C., Chebotar, Y., Hsu, J., Jang. E., Schaal, S., Levine, S., and Brain, G. (2018). cone-connrastive networks: Self-supervised learning from video. In 2018 IEEE international conference on robotics and automation (ICRA), pages 1134 1141. IEEE.
13

---

[57] Silver. D., Huang, A., Maddison, C. 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.. Sulskever, I., Lillierap. T. P. Leach. M., Kavukcuoglu, K., Graepel. T., and Hassabis. D. (2016). Mastering the game of go with deep neural networks and tree search. Nat., 529(7587):484
[58] Sohn, K. (2016). Improved Deep Metric Learning With Multi-Clas N-Pair Loss Objective. In Neural Information Processing Systems, volume 29. Curran Associates, Inc.
[59] Srimivas, A., Laskin, M.. and Abbeel, P. (2020). CURL: contrastive unsupervised representations for reinforcement learning. CoRR, abs/2004.04136
[60] Srivastava, A., Rastogi, A., and et al., A. R. (2023). Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. Truns. Mach. Learn. Res.
[61] Valmeekam, K., Marquez, M., Olmo, A., Sreedharan, S., and Kambhampati, S. (2023). Planbench:
An extensible benchmark for evaluating large language models on planning and reasoning about change. Advances in Neural Information Procensing Systems, 36:38975-3888 -39877.
[62] Venkattaramanujam, S., Crawford, E., Doan, T., and Precup, D. (2019). Self-supervised learning of distance functions for goal- coalcconditioned reinforcement learning. arXiv preprint arXiv: 1907.02996
[63] Wang, T. and Isola. P. (2020). Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In III, H. D. and Singh. A., editors. Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 9929- 9939. PMLR.
[64] Xiao, T., Wang, X., Efros, A. A.. and Darrell, T. (2021). What should not be contrastive in contrastive learning. International Conference on Learning Representations.
[65] Xie, Q. Dai, Z., Du, Y., Hovy, E., and Ncubig, G. (2017). Controllable invariance through adversarial feature learning, Advances in neural information processing systems, 30.
[66] Yin, Y.. Wang. Z., Gu, Y., Huang. H., Chen, W., and Zhou, M. (2024). Relative preference optimization: Enhancing I1m alignment through contrasting responses across identical and diverse prompts. arXiv preprint arXiv:2402. 10958.
[67] Yonetani, R., Taniai, T., Barekatain, M., Nishimura, M., and Kanezaki, A. (2021). Path planning using neural a search. In Meila, M. and Zhang. T., editors, Proceedings of the 38th International Conference on Machine Learning, 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 12039. PMLR.
[68] Zawalski, M., Tyrolski, M., Czechowski, K., Odrzygozdz, T., Stachura, D., Pickos, P., Wu, Y., Kucinski, L., and Milos, P. (2023). Fast and precise: Adjusting planning horizon with adaptive subgoal search. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net.
[69] Zawalski, M., Tyrolski, M., Czechowski, K., Odrzygozdz, T., Stachura, D., Pickos, P., Wu, Y., Lukasz Kucinski, and Milos, P. (2024). Fast and precise: Adjusting planning horizon with adaptive subgoal search.
[70] Zhang, Y., Jiang, H., Miura, Y., Manning, C. D., and Langlotz, C. P. (2022). Contrastive learning of medical visual representations from paired images and text. In Lipton, Z.,, Ranganath, R., Sendak, M., Sjoding, M., and Yeung, S., editors, Proceedings of the 7th Machine Learning for Healthcare Conference, volume 182 of Proceedings of Machine Learning Research, pages 2 25, PMLR.
[71] Zheng, C., Eysenbach, B., Walke, H., Yin, P., Fang. K., Salakhutdinov, R., and Levine, S. (2024). Stabilizing Contrastive RL: Techniques for Offline Goal Reaching. In International Conference on Learning Representations.
[72] Zheng, C., Salakhutdinov, R., and Eysenbach, B. (2023). Contrastive Difference Predictive Coding. In Twelfth International Conference Learning Representations.
[73] Zheng, C., Tuyls, J., Peng, J., and Fysenbach, B. (2025). Can a MISI. fly? analysis and ingredients for mutual information skill learning. In The Thirteenth International Conference on Learning Representations.
14

---

DIGIT JUMPER
3
1
2
3
10
8
2
4
9
15
2
14
(a) N-Puzzle.
(b) Lights Out.
(c) Digit Jumper.
Figure 6: Environments: Our experiments used Sokoban (Fig. 1), the Rubik's Cube (Fig. 5), and the three environments shown above.
A Environments
Sokoban. Sokoban well-known puzzle game in which player pushes boxes onto designated goal positions within confined grid. It is known to be hard from a computational complexity perspective. Solving it requires reasoning over vast number of possible move sequences, making standard benchmark for both classical planning algorithms and modern deep learning approaches [17]. Solving Sokoban requires balancing efficient search with long-term planning. our experiments, we use 12x 12 boards with four boxes.
Rubik's Cube. The Rubik's Cube is 3D combinatorial puzzle with over 4.3 10a possible configurations, making it an ideal testbed for algorithms tackling massive search spaces. Solving the Rubik's Cube requires sophisticated reasoning and planning, as well as the ability to efficiently navigate high- nigh-dimensional state spaces. Recent advances in using neural networks for solving this puzzle, such [1], highlight the potential of deep learning in handling such computationally challenging tasks.
N-Puzzle. N-Puzzle is sliding-tile puzzle with variants such as the 8-puzzle (3x3 grid), 15-puzzle (4x4 grid), and 24-puzzle (5x5 grid). The objective is to rearrange tiles into a predefined order by sliding them into an empty space. It serves classic benchmark for testing the planning and search efficiency of algorithms. The problem's difficulty increases with puzzle size, requiring effective heuristics for solving larger instances.
Lights Out. Lights Out is single-player game invented in 1995. It is grid-based game in which each cell (or light) can be either on off. Pressing cell flips its state and those of its immediate neighbors (above, below, left, and right). Corner and edge lights have fewer neighbors and therefore affect fewer lights. The goal is to press the lights in strategic order to turn off all the lights on the grid.
Digit Jumper. Digit Jumper is grid-based game in which the objective is get from the top-left corner of the board to the bottom-right corner. At each point, the player can move n steps to the left, right, up, or down, where is determined by the number written on the current cell. Digit Jumper is an example of an environment with constant context, as is Sokoban.
B Best-First Search
Best-First Search (BestFS) greedily prioritizes node expansions with the highest heuristic esti- mates, aiming to follow paths that are likely to reach the goal. Although it does not guarantee optimality, BestFS offers a simple and efficient strategy for navigating complex search spaces. The ligh-level pseudocode for BestFS pre- sented in Algorithm 2.
Algorithm Best-First Search [29]
while has nodes to expand do
Take node N with the highest value Select children ni, of N
Compute values U2 for the children Add (n, vy) to the search tree end while
15

---

Grid Search: Hidden Size vs Depth vs Repr Dim (600k training steps)
rеpr dim
repr dim
64
repr dim 128
0.11
0.14
0.03
0.21
0.17
0.14
0.18
0.18
0.15
depth
depth
depth 10
0.15
0.14
0.04
80
0.33
0.16
0.09
0.24
0.17
0.13
10
0.13
0.21
0.09
10
0.19
0.29
0.07
0.27
0.23
0.12
512
1024
2048
512
1024
2048
512
1024
2048
hidden_siie size
bidden_size
hidden size
Figure 7: Grid of network's depth, representation dimension and hidden dimension. The success rate evaluated on cubes scrambled with 10 random moves.
Training Parameters
Loss Parameters
256
0.19
0.33
0.20
0.24
0.40
0.50
0.35
Normalization
0.54
0.28
Batch Size
512
0.20
0.39
0.43
0.19
0.30
Success Rate
0.45
Success Rate
0..55
1024
0.13
0.34
0.28
0.22
0.35
-0.00
1.0
0.28
0.23
0.30
2048
0.15
0.19
0.23
0.25
0.10
1s:05
1e-04
3c-04
1e-03
0.0
1.0
Learning Rate
Excluding the Diagonal
Figure Learning rate batch size grid for Rubik's Cube. The success rate is evaluated on cubes scrambled with 10 random moves after 700k training steps.
Figure 9: CRTR is only effective normalization strategy Sokoban. Effect using negatives in con- trastive learning in Sokoban. We compare the setting where the distance to positives is normalized by the sum over all batch elements only the in-batch nega- tives. The success rate evaluated on cubes scrambled with 10 random moves after 400k training steps.
C Training Details
Code reproduce all results is available in the anonymous repository referenced in the main text. Below, we document the training procedures for the supervised baseline, contrastive baseline, and CRTR.
Training data. For Sokoban, we use trajectories provided by Czechowski et al. [14] and train on dataset of 10 trajectories. For 15-Puzzle, Rubik's Cube, and Lights Out, we generate training trajectories by applying policy that performs random actions, where is set to 150, 21, and 49, respectively. In the case of 15-Puzzle, we additionally remove single-step cycles from the dataset to improve data efficiency. For Digit Jumper, generate training data by sampling random path from the upper-left corner the bottom-right corner on standard 20 20 grid. All grid cells not required for this path are filled by sampling uniformly from the set 1,..., 6. The network for Digit Jumper typically converges after a few hours of training, so we train until convergence observed. For Sokoban, Rubik's Cube, Lights Out, and 1 15-Puzzle, we adopt an unlimited data setup and train all models for two days. This results in the models performing approximately 8 x 10 gradient updates for Rubik's Cube. 7 10 for 15-Puzzih and 10 for Lights Out.
Training hyperparameters. We use the Adam optimizer with constant learning rate throughout training. A learning rate of 0.0003 was found to perform well across all environments, with the exception of Lights Out, where this setting led unstable training. For this environment, we instead reduced learning rate of 0.0001. In all environments,
a batch size of 512. The choice of
16

---

Final Solved Rate (1M training steps)
0.6
Solved Rate
0.4
0.2
0.0
dor backward w/ negs dot backward wal negs dot symmetric w/ negs dot symmetric wo/ negs 12 backward w/ negs 12 backward I wo/ negs
symmetric w/ nogr 12 symmettrcc wo/ negs 12? backward w/ negs 12' hackwardd wa/ negs 12 I symmetric w/ nego 12 symmetric wo/ negs
I1 backward w/ negs
11 backward | wa/ negs
11 symmetric w/ nego
11 symmetric woi negs
Figure 10: Success rate on Rubik's Cube scrambled with 10 random moves, models trained with different contrastive losses. Models using the backward loss consistently achieve hetter performance than those using the symmetric variant. Using the dot product without in-trajectory negatives performs similarly to the & metrie, while combining the dot product with in-trajectory negatives yields the highest success rate. In contrast, combining in-trajectory negatives with symmetric loss results in drop in performance, likely because, in CRTR, such negatives are often closer to the correct solution state-space.
learning rate and batch size was guided by the performance of the contrastive baseline on Rubik's s Cube. Specifically, we evaluated solve rates on cubes shuffled 10 times, as shown in Figure 8. We also conducted grid searches to find the optimal training parameters (learning rate and batch size) for the supervised baseline on Sokoban, Lights Out, and Rubik's s Cube We use the same batch size and learning rate across all methods and environments, with the exception of Lights Out, where increasing the batch size and learning rate the supervised baseline led a higher success rate.
Network architecture. We adopt the network architecture proposed by Nauman al. [45], using 8 layers with a hidden size of 512 and a representation dimension of 64. This configuration was found to yield optimal performance for the contrastive baseline Rubik's 's Cube, illustrated in Figure 7. We observed that this architecture performs well in all environments except for two cases:
* In Sokoban, a convolutional architecture required achieve strong performance.
In Lights Out, the convolutional network was necessary to ensure training stability.
Test set. For Sokoban, we construct a separate test set comprising 100 trajectories, which used to compute evaluation metrics such as accuracy, correlation, and t-SNE visualizations. For all other environments, separate test set unnecessary, as we train for only single epoch. this setting, evaluation is performed directly on unseen data sampled during training.
Contrastive loss. We use the backward version of the contrastive loss, which we found to consis- tently outperform the symmetrized variant on Rubik's 's Cube shown in Figure 10. We also found the backward version to work better on 15-Puzzle and slightly better in the remaining environments.
For Rubik's 's Cube, we use the dot product the similarity metric. Performance across different metrics is presented in Figure 10. While the contrastive baseline performs comparably under the & metric, CRTR achieves significantly better results with the dot product. Based on similar empirical evaluations, we use the following metries for other environments:
Lights Out: & distance,
Digit Jumper and 15- 15-Puzzle: dot product,
17

---

Sokoban: squared &g distance.
We the temperature parameter in the contrastive loss the square root of the representation dimension.
Supervised baseline. The supervised baseline takes as input a pair of states and predicts the distance between them by classifying into discrete bins, where the number of bins corresponds to the maximum trajectory length observed in the dataset.
In all environments, the supervised baseline uses the same architecture the contrastive baseline.
D Evaluation Details
We evaluate all networks on 1000 problem instances per environment. For Rubik's Cube, each instance cube scrambled using 1000 moves. For 15-Puzzle, Lights Out, and Digit Jumper, evaluation boards are sampled randomly. For Sokoban, we follow the same instance generation procedure described by Czechowski al. [14].
E Additional Experiments
A solver. To verify that the improvements achieved by CRTR are not specific to greedy solvers, conducted an additional experiment using the A search algorithm. A employs a heuristic function of the form heuristic a- cost, where varying allows trading off between the search budget required to solve the problem and the average solution length. As shown in Table 1, for the Rubik's s Cube, increasing from 0 (equivalent BestFS) to 500 consistently yields better performance for CRTR compared to CRL. We therefore hypothesize that the improvement reported in Section 5.2 is not specific greedy solvers.
Table 1: CRTR effectiveness is not BestFS specific. A search results the Rubik's 's Cube with node budget of 6000, varying the priority function. CRTR performs better than CRL for all values of o, achicving shorter solution lengths higher solved rates.
0
100
200
300
400
CRTR Avg. Solution Length
56.76
46.35
38.42
32.84
29.16
CRTR Success Rate
0.63
0.62
0.59
0.54
0.33
CRL Avg. Solution Length
62.96
49.88
41.94
36.11
31.77
CRL Success Rate
0.54
0.50
0.44
0.40
0.30
No-search results. The no-search approach selects, at each step, the state that appears most likely to lead toward the solution-based on the learned representation. If the representation were perfect, this strategy would yield optimal solutions. In practice, however, suboptimal representations often cause the agent to wander through latent states far from the goal before eventually converging. As result, the quality of the representation is reflected in the length of these trajectories: the better captures directionality in latent space, the shorter the resulting solutions.
Table 2 reports the average solution lengths for the no-search approach on Rubik's Cube and 15- Puzzle. The results suggest that the representations learned by CRTR are better suited to this approach than those learned by the contrastive baseline, and they significantly outperform those derived from the supervised method. This supports the conclusion that CRTR provides a more reliable notion of direction in latent space. Notably, the average solution lengths for both CRTR and CRL are shorter than the length of training trajectories in 15-Puzzle (150), indicating evidence of trajectory stitching.
We furthermore present the distributions of solution lengths for all the methods in Figure 11.
Accuracy in Sokoban training. During the training of CRL on the Sokoban environment, perfect accuracy acquired almost immediately, due to the method relying on the context, demostrated in Figure 14.
18

---

Rubik
Lights Out
Sokoban
Digit Jumper
15-Puzzle
1.0
Success Rate
0.5
0.0
30
27 Solution Length
61
Solution Length
Solution Length
Solution Length
Solution Length
CRTR without Search CRL without Search
Supervised without Search S Supervised with BestFS
CRTR with BestFS
CRL with BestFS
Figure 11: CRTR produces shorter solutions without explicit search in comparison to baselines. Search can help reduce solution length further. Fraction of boards solved with a solution length of at most æ, comparing CRTR to baselines. Figure 4 in the main text presents analogous results, but only CRTR, for clarity.
Table 2: Average solution length of the baselines and CRTR on Rubik's Cube and 15-Puzzle without using search. Supervised baseline fails to solve Rubik's Cube without search.
Contrastive Baseline
Supervised Baseline
Problem
CRTR
Rubik's 's Cube
448.7
1830.3 119.5
NaN
15 15-puzzle
82.4
1054.3
Digit Jumper analysis. Digit-Jumper is an example of another constant context (defined in Sec. 4. environment, as is Sokoban. It is therefore another environment in which CRL fails rather spectacu- larly and therefore, we observe a similar effect to that seen Sokoban when comparing CRTR to standard CRL. As shown in Figure 12, CRL rapidly achieves 100% training accuracy. However, de- spite this perfect accuracy, the resulting representations exhibit poor correlation with actual temporal structure (Figure 13). This is consistent with the t-SNE visualization (Figure 15): as with Sokoban, CRL collapses each trajectory into a single point in the representation space, discarding temporal information. In contrast, CRTR preserves a clear temporal structure within the latent space (see Figure 15). For non-constant context environments, the difference in representation quality is also visible in success rates, accuracy and correlation, it is however much less pronounced.
Results on SAT. To compare against GNN baselines, we conducted preliminary experiments evaluating CRTR on the 3-SAT task from G4SATBENCH [39] (easy split). When trained on this benchmark, CRTR was unable to achieve nonzero test performance, whereas GNN-based methods reach nearly 90%. We attribute this discrepancy the extremely low-data regime: networks in this benchmark are evaluated after training on only 800 Pemples--rdees of magnitude fewer than those used to train CRTR on tasks where excelled. Supporting this interpretation, we found that simplifying the evaluation by restricting clauses to at most five independent variables (instead of 10-40) substantially improved results. Under this setting, performance increased from random level (15%) to 52% after 5000 training steps, when running the evaluation with the no-search approach. We therefore hypothesize that CRTR is less sample efficient than GNN-based methods, suggesting that further work is needed to improve data efficiency in representation-llamiin approaches.
Results on Sudoku. Given Sudoku's recent popularity as a benchmark for evaluating reasoning in large language models, we also tested our method this task. CRTR was unable to achieve performance above random chance. We attribute this to the goal-conditioned nature of our approach: solving a Sudoku puzzle requires reasoning toward an unknown final configuration, whereas our method assumes access to a well-defined goal state during inference. While this assumption holds for many of the tasks on which CRTR performs well, is violated in Sudoku. To mitigate this, we experimented with adding an artificial "solved state" shared across all trajectories, but this did not substantially improve results. We view this limitation as a promising direction for future work on extending goal-co conditioit architectures to problems with implicit unknown goals.
19

---

jugit.umpprr Accuracy
Digit Jumper Correlation
Sokoban Accuracy
Accuracy
0.0
Correlation
Accuracy
CRTR
CRTR
0.6
0.4
0.4
CRTR
0.3
0.0
0.0
30000
Train Epoch
Evaluation Epoch
Train Epoch
Figure Digit Jumper, CRL
Figure 13: Digit Jumper,
Figure 14: In Sokoban, CRL quickly acquires near-perfect ac- curacy, however this is due to re- lying only superficial features, such walls. Accuracy of clas- sifying whether two states form a positive pair across training: CRTR compared with CRL. The accuracy saturates value smaller than I for CRTR, as result of containing in-trajectory negatives.
quickly acquires near-perfect ac-
CRTR improves temporal struc- ture in robotics euvironments.
curacy, however this is due to re- lying only superficial features
Comparison of Spearman's rank cor- relation metric for CR² (solid) and CRL (dashed) for D4RL, offline
board layout. Accuracy of classifying whether two states form a positive pair across the training, CRTR compared with CRL.
datasets.
Figure 15: CRTR makes representations reflect the structure of the combinatorial task. t-SNE visualization of representations learned by CRTR (left) and CRL (right) for Digit Jumper. Colors correspond to trajectories. CRL representations (right) cluster within trajectories, making them useless for planning.
F Generalization to Temporal Reasoning in Non-Combinatorii Domains
Temporal Structure in D4RL
To investigate whether CRTR also identifies temporal features in non-combinatorial domains, we apply it a dataset of robotic manipulation trajectories (the Adroit dataset from D4RL [23]). Those tasks require using high-dimensional robotic hand to perform fine motor ac- tivities, and are designed to test fine motor control and long-horizon planning. We quantify representation quality
Corvelation (Spearman)
0.8
0.6
0.0
by measuring the predicted distance from each state in trajectory to the final state in trajectory. Specifically, we look the rank correlation between the time step and predicted distance, with a correlation of indicating that the learned representations are highly predictive the temporal distance from each state to the final state.
0.2
Door, CRTR Door, CRL Hammer CRTR Hammer, CRL
Pen. CRTR Pen. CRL
Relocate, CRTR
Relocate, CRL
We look at the correlation through training for CRTR and CRL (Fig. 16). CRTR results in a higher correlation (more than 0.9 in comparison 0.5 0.8 depending on the en- vironment), as well as visibly better training stability for standard CRL, the correlation is visibly unstable through training and in some cases even becomes smaller as the training progresses. This result is little surprising, and not fully clear why does the improvement happen. hypothesize that this because the initial position of the robot differs between trajectories and serves sort of slowly
Figure 16: CRTR improves temporal structure robotics environments. Com- parison of Spearman's 'ss rank correlation met- ric for CR² (solid) and CRL (dashed) for D4RL offline datasets.
20

---

CMI [XX; x |C) vs
CMI [(x; x [C) VS k
Estimated CMI
Estimated CMI
CRTR
CRL
CRTR
CRL
k (number of neighbors)
k (number of neighbors)
Figure 17: CRTR optimizes the conditional mutual information while CRT does not, confirming out theoretical results 4.1. Conditional mutual information estimated in Sokoban (left) and Digit Jumper (right) for representations learned by CRTR and CRL, different values of nearest neighbors used estimation.
changing context, similarly the Rubik's Cube case. We conclude that using CRTR results in better temporal structure in the representation space for non-combinatorial problems.
G Correlation as a Measure of Representation Quality
To assess whether Spearman rank correlation reliable indicator of representation quality, we performed a grid of 96 short runs for each of three environments: Sokoban (12x12), Sokoban (16x16), and the Rubik's ss Cube. We varied four factors: network depth (8, 4, 2), network width (1024, 16), representation dimension (64, 32, 16, 8), and the distance metric used in the contrastive loss (dot product, 2, Eg)
Across all environments, the final Spearman correlation (computed with budget of 1000 nodes) showed a strong relationship with the final success rate: 0.89 for 12x12 Sokoban. 0.80 for 16x16 Sokoban, and 0.90 for the Rubik's ' Cube. These results support the conclusion that Spearman rank correlation a good measure of representation quality.
H
Mutual Information Analysis
To estimate the conditional mutual information, we use NPEET package, which implements the method proposed [37] that uses k-nearest neighbours for entropy estimation. We conduct the analysis using trajectories collected from the Sokoban Digit Jumper environment, utilizing all transitions within these trajectories (> 45k transitions for Sokoban and 20k for Digit Jumper). The variables used in the experiment are defined as follows:
X: Current state embeddings, standardized using z-score normalization (mean 0, standard deviation 1) across the dataset. These embeddings are then projected onto 3-dimensional subspace using Principal Component Analysis (PCA).
X Next state embeddings corresponding transitions from X. The same standardization parameters and PCA transformation applied to X are used for X to ensure consistency.
C: Trajectory identifiers (traj_id) encoded as 2-dimensional vectors sampled from standard bivariate Gaussian distribution (i.e., (O, I2)).
To mitigate the effects of the curse of dimensionality and ensure reliable performance of k-nearest neighbor (kNN)-based estimators, we reduce all high-dimensional representations to low-dimensional spaces (3D for state embeddings. 2D for trajectory identifiers). The conditional mutual information for CRTR and contrastive baseline is reported in Figure 17.
I Ablations
21

---

Repetition Factor Impact
Repetition factor. Our method introduces a single ad- ditional hyperparameter. the repetition factor R. This parameter controls the proportion of in-trajectory nega- tives and is critical for achieving strong performance. As shown in Figure 18, the impact of increasing R varies by environment. For Sokoban, higher values of R lead to only slight decline in performance. In contrast, in many other env environments, excessive repetition can significantly degrade results. While R = 2 is not always optimal, it con- sistently improves performance across all environments we evaluated and serves strong default choice.
Success Rate
0.6
0.3
Rubik (oudger 30001
Sokoban (hudget 3000)
15 Pazzl (budget 200)
0.0
4 Repetition Factor
16
Figure 18: A repetition factor of 2 con- sistently improves the performance. In- creasing the repetition factor for Sokoban, N-Puzzle, and Rubik's ' Cube, respectively.
In Figure 19, we present detailed results showing how varying the repetition factor influences the success rate.
Sokoban
1 15-puzzle
Rubik
Success Rate
0.6
0.8
0.50
Factor
Factur:
Factae
Factor:
Factor:
Facton
0.4
Facter:
Factor:2
0.3
Farton
Pactor:
Pacter: 4
Pacter:
Facton
Pactor:
Facter:
Facton 64
0.0
Factor: 64
0.00
Pactor: 32
6000
1600
2000
4000
60D0
Budget
Budget
Budget
Figure 19: Influence of the repetition factor depends on the environment type. Increasing the repetition factor for Sokoban, N-Puzzle, and Rubik's 's Cube, respectively.
Negatives. We explored alternative methods for incorpo- rating in-trajectory negatives into the contrastive loss. The first approach mimics the standard addition of hard negatives: given a batch B = Ti+ (E..). we sample additional negatives (I;- (ri-)xE(1..BB) and compute the loss as
(f (;-),,;;))) (xx((((),
==
Σ log Ejft exp(((2,))))) +
We considered three strategies for selecting in-trajectory negatives: sampling state uniformly at random, choosing the first state, or choosing the last state of the trajectory. For Rubik's Cube, instead of choosing the state- whech identical for all trajectories- we sample a random state farther from the solution to serve negative.
10 Shuffles
15 Shuffles
20 Shuffles
Success Rate
Success Rate
Success Rate
8.32
0.36
0.3
0.04
Normili random Norpalie aeddddd zervee
0.0
CRTR CRL
L long
Ungormal tzed leng random
poraaazzzdd yeed zerons Unnormahizdd
Normelized l long
acc.
CRTR CRL
Norpalia sedl I zervin
Unnormal lzzd losg random Unnormalin Unnormahizdd yea zerons
random
Normalized annd lnng
CRTR CRL
Normali sealizdd I zerune
Unnormal Lzed long random
random
Unnermall Unncrmalized Lea zeroos
Figure 20: Rubik's Cube, CRTR outperforms all negative sampling strategies, when number of scrambles increases. Comparison of different methods for introducing in-trajectory negatives in the Rubik's Cube environment, with an increasing number of cube scrambles. While normalized negatives perform similarly CRTR for small number of scrambles, their performance deteriorates the number of scrambles increases.
22

---

Solved Rate per Run (after 400k training steps)
Solved Rate
0.30
0.15
0.00
CRTR
Normalized first
CRL
Normalized random
last
Unnormalized first
Unnormalized random
last
Normalized |
Unnormatized
|
Figure 21:
We compare various methods for introducing in-trajectory negatives the Sokoban environment and find that only CRTR yields effective results.
As shown in Figures 20 and 21, training with this approach did not yield strong performance. We hypothesized that the large prediction error introduced by the in-trajectory negatives (xx-) caused excessively large gradients, destabilizing training. To mitigate this, we applied normalization scheme: ensuring that the vector [f(21,11-) the B x B matrix
((BB, has the same Frobenius norm as
[(x1, ((11,02))
f(F8, f(IB,
f(zg,TB+)]
This normalization enabled achieving comparable performance to CRTR on Rubik's Cube scrambled 10 times (Figure 20). However, CRTR still outperforms all negative sampling strategies cubes scrambled 15 and 20 times.
For Sokoban, the only approach that consistently improved performance is CRTR, as demonstrated in Figure 21. We hypothesize that this because removing contextual information more challenging in Sokoban than in Rubik's Cube. In the latter, the context is more local and changes gradually over time, making softer, while the context in Sokoban constant throughout trajectory. This discussed in detail in Section 4.1.
While at first glance, repeating trajectories in batch may seem equivalent to sampling in-trajectory hard negatives, the two approaches are different. In standard contrastive learning (as in CRL), an anchor æ pulls its positive r. closer and pushes negatives (e.g., y+) away. However, negatives like 3+ are simultaneously pulled by their own anchors (e.g.. y), which limits how far they are pushed by x. In contrast, when using in- in-trajectory negatives without anchoring them (e.g., pushes x away, but x_ has no anchor). these states can drift arbitrarily far in representation space. This is problematic, especially since in-trajectory negatives are harder (closer in structure), which results in stronger gradient updates. Our proposed method, CRTR, addresses this by anchoring all in-trajectory negatives. This keeps trajectories coherent and prevents such drift.
J Computational Resources
All training experiments were conducted using NVIDIA A100 GPUs and took between 5 and 48 hours each. The solving runs ranged from 10 minutes 10 hours. In total, the project required approximately 30,000 GPU hours to complete.
K Things We Tried That Did Not Work
Using separate encoders for future and present states did not improve performance.
23

---

Adding extra layers to encode the action led to lower success rates.
Using only in-trajectory negatives degraded performance.
Modifying how current states are sampled in CRL (e.g., deviating from uniform sampling) did not yield improvements.
Using A solver with our representations could be greatly improved. Because distances in the latent space are only monotonically correlated not linearly correlated -with actual distances, a modification to A" that would account for these discrepancies could bring huge gains.
Distances between Rubik's Cube states, measured by the number of actions, almost always satisfy the triangle inequality with equality. Consequently, this metric cannot be faithfully embedded in Euclidean space, where equality in the triangle inequality occurs only for collinear points.
Since Rubik's Cube actions are not commutative, a faithful Cayley graph structure could only emerge in Euclidean space where vector addition is aancomutattve--whie would require a highly non-standard space.
24

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: The claims made match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
It fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The paper discusses the limitations of the work.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "LLimitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
25

---

Answer: [Yes]
Justification: Paper provides the full set of assumptions and complete proof for the one theoretical claim it makes.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
The paper fully discloses all the information needed to reproduce the main results. The code is provided well as scripts needed for data generation.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
26

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
The paper provides open access the data and code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co
wwm.iiisssiiiiiioooio for more details. http://gguidee/ooeeeuee
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: CCodeSubdeeeaeeee
cttppssccppbbcccccepeecepcpccccc / ssi onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification:
The paper does specify the training and test details, partially in Section 5 and partially in Appendices C and
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: In Figure 2, we present the results obtained from three random seeds. With one exception. the training exhibits consistently stable performance, which supports the reliability of the remaining plots.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
27

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide the details in Appendix J.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: /nwuurius.... cc cccpppplcc/////eceeeceeecccllllt csQuidelines!
Answer: [Yes]
Justification:
The research conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The paper states that the socictal impacts are the standard socictal impacts resulting from developing machine learning models, which is true.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
28

---

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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
The creators and original owners of all assets used are properly credited.
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 dataset.
The authors should state which version of the asset used and, possible, include 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.
29

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswitie codee com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We release the code which is well documented.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
30

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
The core method development does not involve LLMs.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: poetpa:/an /ttpr//ees.p///// cc/Conferences//22222 /LLM) for what should or should be described.
31


Reviews:
[{"review_id": "5dO9VdE2uS", "reviewer": "Reviewer_3TVa", "summary": "The paper presents CR$^2$ (Contrastive Representations for Combinatorial Reasoning), which extends temporal contrastive learning with in-trajectory negatives to learn embeddings that double as effective search heuristics. \nBy revealing how standard contrastive RL collapses onto static context and then maximizing conditional mutual information via a simple “repetition factor,” the method preserves the temporal signal crucial for planning. Evaluated on several tasks like complex puzzles and control problems, CR2’s embeddings align closely with true state-space distances, lifting correlation metrics and solve rates and, in many settings, matching the performance of heavyweight search-centric systems. However, further analysis confirms that explicit search is still required on the harder instances.", "questions": "None", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**Strenghts**\n\n- The paper is well written and easy to follow. \n- The paper identifies an important flaw with standard contrastive RL approaches, especially for combinatorial reasoning tasks, and provides a solution -- using CR$^2$ that learns to ignore the context and focus on the temporal and causal structure. \n- The paper provides further analysis and experiments where they investigate CR$^2$ such as testing if search is still needed, experimenting in cases where the context is dynamic. \n\n**Weaknesses**\n\n- The conclusions from the t-SNE plots are not very clear, and the claim that the plots for CRL form clusters, indicating they learn the context (walls), while CR$^2$ seems to be a bit hand-wavy. Perhaps this argument could be strengthened.\n\n\nOverall, I think the paper makes a good contribution, and I recommend acceptance", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for their effort to assess our work! The main concern raised appears to be that our conclusions drawn from the t-SNE plots require stronger empirical support. In order to do that we propose two additional experiments, one based on training a linear classifier on top of our representations and another one using a mutual information estimator, based on our learned representations. The details are contained in the response below.\n\n**> The conclusions from the t-SNE plots are not very clear, and the claim that the plots for CRL form clusters, indicating they learn the context (walls), while CR$^2$ does not, seems to be a bit hand-wavy. Perhaps this argument could be strengthened.**\n\nTo support the conclusions drawn from the t-SNE plots, we conducted two experiments to assess how much information about the context do the CRL and CR$^2$ representations contain.\n\n# Experiment 1\nThe experimental setup is simple: we train a linear classifier using representions from CRL, CR$^2$, or a randomly initialized network as input and use the trajectory ID as the prediction target. The trajectory ID serves as a proxy for the static context (e.g., wall configuration), since each trajectory is tied to a fixed environment layout. We use Sokoban trajectories (>49k transitions), with 70% of transitions from each trajectory used for training and the remaining 30% for testing. To obtain reliable accuracy estimates, we perform 6-fold cross-validation.\n\nIntuitively, if the representations contain contextual information, it should be easier for the classifier to predict the correct trajectory ID.\n\nThe results are as follows (accuracy mean ± standard deviation):\n- CRL: 88.16 ± 0.30%,\n- CR$^2$: 1.80 ± 0.09%,\n- Randomly initialized representations: 3.32 ± 0.15%.\n\nThe results show that trajectory ID—which reflects contextual information—can be accurately predicted from CRL representations in Sokoban. CR$^2$ representations achieve accuracy significantly lower than a random baseline, indicating that contextual signals are effectively suppressed.\n\n# Experiment 2\nIn the following experiment, we estimate the mentioned mutual information using [NPEET package](https://github.com/gregversteeg/NPEET), which implements the method proposed in [1] that uses k-nearest neighbours to get entropy estimations. We conduct the analysis using trajectories collected from the Sokoban or Digit Jumper environment, utilizing all transitions within these trajectories (>45k transitions for Sokoban and >20k for Digit Jumper). The variables used in the experiment are defined as follows:\n- $X$: Current state embeddings, standardized using z-score normalization (mean 0, standard deviation 1) across the dataset. These embeddings are then projected onto a 3-dimensional subspace using Principal Component Analysis (PCA).\n- $X^+$: Next state embeddings corresponding to transitions from $X$. The same standardization parameters and PCA transformation applied to $X$ are used for $X^+$ to ensure consistency.\n- $C$: Trajectory identifiers (traj_id) encoded as 2-dimensional vectors sampled from a standard bivariate Gaussian distribution (i.e., $\\mathcal{N}(0, I_2)$).\n\nNote: To mitigate the effects of the curse of dimensionality and ensure reliable performance of k-nearest neighbor (kNN)-based estimators, we reduce all high-dimensional representations to low-dimensional spaces (3D for state embeddings, 2D for trajectory identifiers).\nIn the tables below, we report $I(X;X^+ \\mid C)$, $I(X;X^+)$, and $I(X^+;C)$ for both CRL and CR$^2$ using 3 different numbers of neighbours ($k=3,4,5$) used for mutual information estimation. We report 3 different values of $k$ to demonstrate the stability of the estimation.\n## Sokoban\n### CRL Mutual Information (MI)\n| MI type     |   $k=3$ |  $ k=4$ |  $ k=5 $|\n|-------------------|-------|-------|-------|\n| $I(X;X^+ \\mid C)$ |  1.2  |  1.13 |  1.07 |\n| $I(X;X^+)$     | 11.01 | 10.93 | 10.87 |\n| $I(X^+;C)$     |  9.8  |  9.8  |  9.8  |\n\n### CR$^2$ Mutual Information (MI)\n| MI type     |   $k=3$ |  $ k=4$ |  $ k=5 $|\n|-------------------|-------|-------|-------|\n| $I(X;X^+ \\mid C)$ |  6.47 |  6.56 |  6.6  |\n| $I(X;X^+)$     |  7.26 |  7.2  |  7.14 |\n| $I(X^+;C)$   |  0.79 |  0.64 |  0.55 |\n\n## Digit Jumper\n### CRL Mutual Information (MI)\n| MI type     |   $k=3$ |  $ k=4$ |  $ k=5 $|\n|-------------------|-------|-------|-------|\n| $I(X;X^+ \\mid C)$ |  1.76 |  1.97 |  2.14 |\n|  $I(X;X^+)$     |  9.4  |  9.41 |  9.4  |\n| $I(X^+;C)$     |  7.64 |  7.44 |  7.26 |\n\n### CR$^2$ Mutual Information (MI)\n| MI type     |   $k=3$ |  $ k=4$ |  $ k=5 $|\n|-------------------|-------|-------|-------|\n| $I(X;X^+ \\mid C)$ |  3.96 |  4.12 |  4.24 |\n|  $I(X;X^+)$     |  5.84 |  5.8  |  5.75 |\n|$I(X^+;C)$   |  1.89 |  1.68 |  1.51 |\n\nWe observe that CR$^2$ indeed minimizes $I(X^+; C)$, while maximizing $I(X;X^+)$, which leads to the desired maximization of $I(X;X^+\\mid C)$ for both environments. In contrast, CRL baseline results in high values of $I(X^+;C)$, which contribute to relatively low conditional mutual information $I(X;X^+\\mid C)$.\n\nWe have briefly described this analysis at the end of Section 4.2 (Learning Representations that Ignore Context) and put the detailed results in the Appendix.\n\n**> Line 210 -- Instead of Algorithm 4.3, it should say Figure 3.**\n\nThank you for catching this. We have corrected the reference on Line 210 to \"Figure 3\" as suggested.\n\n**Summary:** \n\nWe hope these additions meaningfully strengthen the paper and address the reviewer’s concerns. We would be happy to engage further or provide additional experiments if needed.\n\n**References:**\n\n[1] Kraskov, A., et al., Phys. Rev. E, Estimating Mutual Information"}, {"review_id": "dAVwrKRyHS", "reviewer": "Reviewer_TTPY", "summary": "The paper introduces CR2, a contrastive representation learning method for solving multi-step combinatorial problems. The paper demonstrates the failure of standard contrastive learning approaches in solving combinatorial problems due to the presence of excessive non-essential information, commonly referred to as \"common context,\" in all instances of the same trajectory. Based on this observation, the paper goes on to propose a method to introduce in-trajectory negatives. This is achieved simply by replicating the trajectories a fixed number of times in the same batch and leaving the rest of the CL pipeline as is, which has the effect of randomly introducing negatives from the same trajectory. The experiments demonstrate that CR2 improves significantly over the naive CRL objective. It is also competitive with supervised value-based methods.", "questions": "1. Isn't repeating trajectories (as described in Figure 3) in a batch equivalent to randomly (perhaps not uniformly randomly) sampling in-trajectory \"hard negatives\"? \n\n2. Why does the training objective saturate for CR2 in Figure 5? Is it to be expected due to in-trajectory negatives?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Quality**: All claims made in the paper are well-supported by targeted experiments. The limitations of the current work, i.e., the use of deterministic known dynamics and noiseless states, are clearly stated. \n\n**Clarity**: The paper is well written and easy to follow. The intuition provided using probabilistic arguments is helpful in understanding the general principle behind the method in a task-agnostic manner.\n\n**Significance**: Using deterministic dynamics and noiseless states makes the experimental setup seem somewhat toy-like. The paper's impact could be greatly enhanced if experiments on more realistic tasks were also included.\n\n**Originality**: Although the idea in the paper closely relates to the concept of \"hard negatives\" in general contrastive representation learning literature, I believe the way the idea is presented and instantiated is original.", "was_revised": true, "final_justification": "I maintain my score. The paper presents a novel idea that works, but the experiments, which only include tasks with deterministic dynamics and noiseless states, are a bit toyish.", "rebuttal": "We thank the reviewer for reading our paper and providing feedback! \n\n**> Isn't repeating trajectories (as described in Figure 3) in a batch equivalent to randomly (perhaps not uniformly randomly) sampling in-trajectory \"hard negatives\"?**\n\nThank you for the thoughtful question. While at first glance repeating trajectories in a batch may seem equivalent to sampling in-trajectory hard negatives, the two approaches behave differently. In standard contrastive learning (as in CRL), an anchor $x$ pulls its positive $x_+$ closer and pushes negatives (e.g., $y_+$) away. However, negatives like $y_+$ are simultaneously pulled by their own anchors (e.g., $y$), which limits how far they are pushed by $x$.\nIn contrast, when using in-trajectory negatives without anchoring them (e.g., $x$ pushes $x_-$ away, but $x_-$ has no anchor), these states can drift arbitrarily far in representation space. This is problematic, especially since in-trajectory negatives are harder (closer in structure), which results in stronger gradient updates.\nOur proposed method, $CR^2$, addresses this by anchoring all in-trajectory negatives. This keeps trajectories coherent and prevents such drift. Empirical evidence supporting this can be found in Appendix G, Figures 23 and 24. \nIn the camera-ready version, we will include this discussion in section 4.3 (Method used in practice) \n\n**> Why does the training objective saturate for CR2 in Figure 5? Is it to be expected due to in-trajectory negatives?**\n\nYes, this saturation is expected and arises due to the use of in-trajectory negatives. Figure 5 shows, for each anchor, how often the closest state is its true positive. With in-trajectory negatives, it’s possible for a negative from the same trajectory to be embedded closer than the positive. Additionally, once wall-based features are no longer sufficient for classification, the model may start relying on more abstract or subtle features. As a result, positives from other trajectories may occasionally be embedded closer to an anchor, contributing to the observed saturation.\nWe have included this explanation in Section 5.3, where the Figure 5 is described.\n\n**> Using deterministic dynamics and noiseless states makes the experimental setup seem somewhat toy-like. The paper's impact could be greatly enhanced if experiments on more realistic tasks were also included.**\n\nThank you for this suggestion. We agree that extending the method to more realistic tasks is an important direction. In this work, we focused on deterministic dynamics and noiseless states, as these are standard in combinatorial puzzles. As discussed in the Limitations section, we plan to extend our approach to real-world combinatorial problems, such as robotic assembly and retrosynthesis. If the reviewer has specific suggestions for additional practical combinatorial domains, we would be happy to include them in the Limitations section for the camera-ready version and explore them in future work."}, {"review_id": "KMWM2v3Ofm", "reviewer": "Reviewer_g7ur", "summary": "Contrastive learning fails on combinatorial reasoning tasks because it learns to exploit superficial contextual features that disregard relevant temporal information (Fig. 1). The failure occurs when there's a constant \"context\" (like wall patterns in Sokoban) that makes it trivial to distinguish between trajectories. Standard contrastive learning samples negatives from different trajectories, so models can achieve perfect accuracy by identifying the context rather than learning meaningful temporal relationships (Fig. 5-6). To resolve this, CR^2 samples negative examples (temporally distant states) from within the same trajectory (in-trajectory negatives) rather than only from different trajectories. \n\nCR^2 optimises $I(X; X_+) - I(X_+; C)$, encouraging representations that capture temporal structure while being invariant to context. Practically, this is done via a simple modification (Fig. 3) to how the data is loads; repeat trajectory indices in batches with a \"repetition factor\" (typically 2). Evaluation across 5 different settings (Sokoban, Rubik's Cube, N-Puzzle, Lights Out, Digit Jumper) shows that CR² consistently outperforms standard contrastive learning (Fig. 4) and matches supervised baselines using Best-First Search planning.", "questions": "1.\tIn the setting with partial observability, stochastic dynamics – Assumption 4.2 would not hold. How sensitive is the method to these violations?  Can we characterise exactly which MDP structures satisfy this assumption?\n2.\tWere other search algorithms evaluated apart from best-first search – that weren’t included in the manuscript? Example, A* ?\n3.\tWe can see that the training objective minimises I(X⁺; C) i.e., mutual information between representations and context, should lead to context invariance. However, this is not shown empirically. Therefore, could we see how much the mutual information decreases during training for the different settings.\n\nWe can see that search is necessary to get appropriate solutions -- with certain environments requiring more of the search budget (6000 nodes). \n\n4.\tCan we quantify the relationship between representation quality (measured by correlation with true distances) and required search budget? \n5.\tDoes a 10% improvement in correlation translate to a predictable reduction in nodes explored? \n6.\tHow do these trade-offs change with problem size (15-puzzle vs. 24-puzzle)?\n7.\tIs this relationship linear, logarithmic, or domain-dependent?\n8.\tIs there a threshold correlation value below which search becomes ineffective?\n\n\n9.\tFor Fig.8 would it be possible to include confidence intervals? this would allow a concrete evaluation of the CRL and CR^2 and importantly the robustness of the evaluation. \n10.\tThe results show that having higher repetition factor degrade performance. Is this because the higher R is increasing too much noise in the training signal?\n11.\tMinor point: Last equality of equation 4 is repeated. It would be good to remove it?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strength: Overall, I like this paper – it provides a nice example of why standard contrastive learning fail on combinatorial problems, proposes a simple fix using in-trajectory negatives and shows consistent improvements across multiple domains. Improves search efficiency when combined with Best-First Search planning. The code is made public for future extensions. \n\nLimitation:  Despite the well-structured evaluations, there are some limitations. The problems considered are quite small-scale and evaluation on more realistic problem settings might be useful to evaluate the significance of the contribution. Some of the results presented seem to be for limited (i.e., 3) or single seeds, making it difficult to evaluate their robustness. The restrictive Assumption 4.2 (X_i ⊥ C | X_j) may not hold in many real-world domains with partial observability or stochastic dynamics, limiting applicability beyond clean combinatorial puzzles. Furthermore, the paper shows that despite the improvements in representation, CR^2 still cannot replace search entirely – particularly in limited budget settings. Table 1 shows solve rates near zero for complex problems without search (15-puzzle). This suggests we may be hitting fundamental limits of what these types of representations can achieve for combinatorial reasoning.", "was_revised": true, "final_justification": "All my concerns were addressed during the review process - I recommend acceptance.", "rebuttal": "We thank the reviewer for their detailed and constructive feedback.\n\n**> Were other search algorithms evaluated apart from best-first search – that weren’t included in the manuscript? Example, A\\*?**\n\nFollowing the reviewer’s comment, we ran preliminary A* experiments on the Rubik’s Cube to test whether CR$^2$ representations remain effective beyond greedy search. Using a node budget of 6000 and varying $\\alpha$ in the priority function $heuristic + \\alpha \\cdot cost$ for $\\alpha \\in {0, \\dots, 400}$, CR$^2$ consistently outperforms CRL in both success rate and solution length. This suggests that CR$^2$ generalizes well across search strategies. We will clarify this point and include the results in the revised paper. The full results are in the table below:\n| $\\alpha$              \t| 0 \t| 100   | 200   | 300   | 400   |\n|-----------------------------|-------|-------|-------|-------|-------|\n| **CR$^2$ Avg. Solution Length** | 56.76 | 46.35 | 38.42 | 32.84 | 29.16 |\n| **CR$^2$ Success Rate**     \t| 0.63  | 0.62  | 0.59  | 0.54  | 0.33  |\n| **CRL Avg. Solution Length**\t| 62.96 | 49.88 | 41.94 | 36.11 | 31.77 |\n| **CRL Success Rate**       \t| 0.54  | 0.50  | 0.44  | 0.40  | 0.30  |\n\n\n\n\n**> We can see that the training objective minimises I(X⁺; C) i.e., mutual information between representations and context, should lead to context invariance. Therefore, could we see how much the mutual information decreases during training for the different settings.**\n\nWe thank the reviewer for the suggestion. To empirically validate that the training objective minimizes $I(X^+; C)$ while maximizing $I(X; X^+ \\mid C)$, we estimate mutual information using the NPEET package [3], which uses a k-nearest neighbor method. We analyze over 45k transitions from Sokoban and over 20k from Digit Jumper.\n\nThe variables are defined as:\n- $X$: Z-score normalized current state embeddings,\n- $X^+$: Next state embeddings, standardized and projected using the same parameters as $X$,\n- $C$: Trajectory IDs, encoded as 2D samples from $\\mathcal{N}(0, I_2)$.\n\nTo mitigate the curse of dimensionality, we reduce all variables to low-dimensional spaces (3D for embeddings, 2D for $C$). The tables report $I(X; X^+ \\mid C)$, $I(X; X^+)$, and $I(X^+; C)$ for CRL and CR$^2$ using $k = 3, 4, 5$ to show estimation stability.\n## Sokoban\n### CRL\n| MI type     |   $k=3$ |  $ k=4$ |  $ k=5 $|\n|-------------------|-------|-------|-------|\n| $I(X;X^+ \\mid C)$ |  1.2  |  1.13 |  1.07 |\n| $I(X;X^+)$     | 11.01 | 10.93 | 10.87 |\n| $I(X^+;C)$     |  9.8  |  9.8  |  9.8  |\n\n### CR$^2$ \n| MI type     |   $k=3$ |  $ k=4$ |  $ k=5 $|\n|-------------------|-------|-------|-------|\n| $I(X;X^+ \\mid C)$ |  6.47 |  6.56 |  6.6  |\n| $I(X;X^+)$     |  7.26 |  7.2  |  7.14 |\n| $I(X^+;C)$   |  0.79 |  0.64 |  0.55 |\n\n## Digit Jumper\n### CRL \n| MI type     |   $k=3$ |  $ k=4$ |  $ k=5 $|\n|-------------------|-------|-------|-------|\n| $I(X;X^+ \\mid C)$ |  1.76 |  1.97 |  2.14 |\n|  $I(X;X^+)$     |  9.4  |  9.41 |  9.4  |\n| $I(X^+;C)$     |  7.64 |  7.44 |  7.26 |\n\n### CR$^2$\n| MI type     |   $k=3$ |  $ k=4$ |  $ k=5 $|\n|-------------------|-------|-------|-------|\n| $I(X;X^+ \\mid C)$ |  3.96 |  4.12 |  4.24 |\n|  $I(X;X^+)$     |  5.84 |  5.8  |  5.75 |\n|$I(X^+;C)$   |  1.89 |  1.68 |  1.51 |\n\nWe observe that CR$^2$ indeed minimizes $I(X^+; C)$, while maximizing $I(X;X^+)$, which leads to the desired maximization of $I(X;X^+\\mid C)$ for both environments. In contrast, CRL baseline results in high values of $I(X^+;C)$, which contribute to relatively low conditional mutual information $I(X;X^+\\mid C)$.\n\nWe have briefly described this analysis at the end of Section 4.2 (Learning Representations that Ignore Context) and put the detailed results in the Appendix.\n\n**> Can we quantify the relationship between representation quality (measured by correlation with true distances) and required search budget?**\n\nThis is a highly relevant and insightful question. While a full analysis is ongoing, we investigated a closely related quantity: the relationship between correlation and the success rate at a fixed budget (6000 nodes).\n\nWe consider three environments: Sokoban with 12x12 boards, Sokoban with 16x16 boards and the Rubik’s Cube. We run a grid of 96 short trainings for each one of them, varying the network architecture hyperparameters and the metric choice for contrastive learning.\n\nWe find that for each one of them, a function of the form $\\frac{1}{a(x - b)^2} +c$ describes this relationship well. The performance is almost constant up to some point and for values close to 1 grows very quickly.\n\nWe find the optimal values to be:\n- Sokoban 12x12 $a=5000, b=0.05, c=0.05$ with MSE 0.02,\n- Sokoban 16x16 $a = 800, b = 0, c=0.05$ with MSE 0.14,\n- Rubik’s Cube $a=1000, b=0.02, c=0$ with MSE 0.02.\n\nTherefore, the solved rate is effectively 0 for correlation smaller than:\n- Sokoban 12x12: $0.8$,\n- Sokoban 16x16: $0.85$,\n- Rubik’s Cube: $0.9$. \n\n\n**> For Fig.8 would it be possible to include confidence intervals?**\n\nThank you for the suggestion. During the rebuttal period, we re-ran the experiments from Figure 8 using 5 random seeds, varying both the train-test split and data order. The standard error was consistently low across tasks, with the exception of CRL on the Hammer and Pen tasks, where it reached approximately 0.05 still not affecting the overall conclusions. We will update the figure in the revised version to include confidence intervals.\n\n\n\n**> The results show that having higher repetition factor degrade performance. Is this because the higher R is increasing too much noise in the training signal?**\n\nIncreased noise in the training signal is indeed a likely contributor to reduced performance. While more negative examples typically reduce gradient variance [1], a high repetition factor can counteract this by introducing highly similar negatives, limiting this benefit. Additionally, higher repetition may introduce bias, as observed in small-batch contrastive learning setups [2]. We have added this discussion to Section 5.6.\n \n**> Minor point: Last equality of equation 4 is repeated. It would be good to remove it?**\n\nThank you for identifying this typo. It has been corrected in the revised version.\n\n**> The problems considered are quite small-scale and evaluation on more realistic problem settings might be useful to evaluate the significance of the contribution.** \n\nThank you for the comment. We agree that evaluating on real-world tasks such as retrosynthesis or robotic assembly would further demonstrate the practical significance of our method, and we note this direction in the Limitations section. However, we would like to clarify that the problems we study are not small-scale. For instance, the Rubik’s Cube has a state space of $4.3 \\times 10^{19}$, and to our knowledge, no existing method has learned high-quality representations in this domain. Successfully doing so highlights our method’s ability to model highly complex and nontrivial distributions.\n\n**> Some of the results presented seem to be for limited (i.e., 3) or single seeds, making it difficult to evaluate their robustness.**\n\nThank you for pointing this out. Some of the original plots were indeed generated from a limited number of seeds. We are currently re-running all experiments with 5 random seeds to improve statistical robustness. Several runs have already completed, and the results closely match those previously reported. We will include updated figures with error bars in the camera-ready version.\n\n**> In the setting with partial observability, stochastic dynamics – Assumption 4.2 would not hold. How sensitive is the method to these violations? Can we characterise exactly which MDP structures satisfy this assumption?**\n\nThank you for this insightful question. Assumption 4.2 requires the context to remain fixed and fully observable across a trajectory, which may not hold in partially observable environments. However, such settings are beyond our scope. We focus on fully observable, deterministic environments—an important and challenging class that includes combinatorial puzzles (e.g., Sokoban, 15-puzzle, Rubik’s Cube) and practical problems (e.g., Travelling Salesman, Robotic Assembly, Retrosynthesis), and an active area of research [3, 4, 5]. We have clarified this in the Limitations section.\n\n**> Furthermore, the paper shows that despite the improvements in representation. This suggests we may be hitting fundamental limits of what these types of representations can achieve for combinatorial reasoning.**\n\nThank you for this observation. We agree that understanding the limits of learned representations is important and underexplored, and see our work as a step in that direction. While success rates on complex tasks like the 15-puzzle remain low under strict step budgets without search, Appendix E shows that relaxing episode length substantially improves solve rates. This suggests the limitation stems more from fixed planning horizons than a lack of generalization. Additionally, on the 15-puzzle, both CR$^2$ and CRL produce solutions shorter than those in the demonstration set, indicating meaningful generalization. CR$^2$ also significantly outperforms CRL on both the 15-puzzle and Rubik’s Cube, highlighting its strength even without search.\n\n\n**References:**\n\n[1] Cho, J., et al.,  TMLR 2024, Mini-Batch Optimization of Contrastive Loss\n\n[2] Chen. C., et al., NeurIPS 2022, Why do We Need Large Batchsizes in Contrastive Learning? A Gradient-Bias Perspective\n\n[3] Kraskov, A., et al., Phys. Rev. E, Estimating Mutual Information\n\n[4] Lehnert, L., et al., COLM 2024, Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping\n\n[5] Bush, T., et al., ICLR 2025, Interpreting Emergent Planning in Model-Free Reinforcement Learning\n\n[6] Zawalski, M., et al., ICLR 2023, Fast and Precise: Adjusting Planning Horizon with Adaptive Subgoal Search"}, {"review_id": "kIgO1hBOHb", "reviewer": "Reviewer_BU6x", "summary": "This paper introduces CR² (Contrastive Representations for Combinatorial Reasoning) to address limitations of standard contrastive learning (CRL) when applied to combinatorial problems. The authors identify that CRL fails in combinatorial domains because it learns to exploit superficial context features (like wall layouts in Sokoban) rather than temporal structure. CR² addresses this by incorporating \"in-trajectory negatives\" - negative samples from the same trajectory rather than just across different trajectories. By including negatives that have similar context (similar wall layouts), the model is forced to learn temporal structure instead of just memorizing context. This is done by repeating trajectories in the batch, so there is a higher chance that negatives come from trajectories with similar context.\n\n\nMain Claims of Contributions:\n\n- Failure Mode Identification: Standard CRL fails on combinatorial problems by learning context-dependent representations that ignore temporal structure\n- Novel Algorithm: CR² optimizes I(X; X⁺) - I(X⁺; C), encouraging representations that maximize temporal mutual information while minimizing context dependence", "questions": "Questions:\n\n- The \"repetition factor\" hack seems unprincipled - why this specific approach?\n- All experiments use BestFS. How do the learned representations perform with other search algorithms (A*, MCTS, etc.)? Is the improvement specific to greedy search strategies?\n- The definition of context (Definition 4.1) works well for Sokoban where static elements are clearly separable, but how would you formally define context in domains like Rubik's Cube where you mention \"pseudo-context\"? Could you provide a more general mathematical framework for identifying context across different combinatorial domains?\n- The conditional independence assumption is crucial to your theoretical analysis. Have you empirically validated this assumption across your test domains? What happens to your theoretical guarantees when this assumption is violated?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n- The paper addresses an important problem in representation learning: capturing temporal dependencies to enable better planning and improve sample efficiency. The combination of contrastive learning with novel approaches shows promise for learning effective representations.\n\nWeaknesses:\n\nThis work requires substantial development (sufficient novelty and additional contributions) before publication. The main issues are:\n\n \n- The core insight about negative sample selection is reasonable but offers minimal novelty. The proposed solution appears to be an ad-hoc fix rather than a principled algorithm, with weak connections between the theoretical objective and actual implementation.\n\n- The method makes strong, unrealistic assumptions. It only works if we can cleanly separate \"context\" from \"temporal dynamics\" - a strong assumption that breaks down in many domains. The paper provides no systematic approach for identifying context in practice.\n\n- Inadequate Experimental Validation: The evaluation doesn't support the broad claims made. Despite the title suggesting general \"combinatorial reasoning,\" testing is limited to puzzle games with no evidence of effectiveness on real combinatorial problems (traveling salesman, vehicle routing, etc.). Additional issues include: No results across multiple random seeds, insufficient analysis of learned representations beyond basic correlation, and missing comparisons to state-of-the-art methods (Graph Neural Networks, neural combinatorial optimization, Monte Carlo Tree Search variants).\n\n\n- Fig. 3 highlights the difference in the code for CRL and CR2 (which is unusual to add in the main paper). This again suggests an engineering modification rather than a research contribution. \n\n- The related work section is superficial and fails to clearly position the work within existing literature.\n\nOverall Assessment:\n\nThe work appears more like an engineering hack than a novel research contribution and needs significant theoretical development and experimental validation before being ready for publication.", "was_revised": true, "final_justification": "I appreciate the author response and have read all the other reviews. Most of my gaps in understanding and concerns have been answered. I am able to appreciate the novelty regarding leveraging trajectory structure to construct informative negatives when using contrastive learning to solve combinatorial problems. The new experiments and analysis strengthen the paper. I have increased the score.", "rebuttal": "We thank the reviewer for the detailed feedback, which we believe will strengthen the paper. \n\n**> All experiments use BestFS. How do the learned representations perform with other search algorithms (A\\*, MCTS, etc.)? Is the improvement specific to greedy search strategies?**\n\nFollowing the reviewer’s comment, we ran preliminary A* experiments on the Rubik’s Cube to test whether CR$^2$ representations remain effective beyond greedy search. Using a node budget of 6000 and varying $\\alpha$ in the priority function $heuristic + \\alpha\\cdot cost$ for $\\alpha \\in {0, \\dots, 400}$, CR$^2$ consistently outperforms CRL in both success rate and solution length. This suggests that CR$^2$ generalizes well across search strategies. We will clarify this point and include the results in the revised paper. The full results are in the table below.\n| Coefficient              \t| 0 \t| 100   | 200   | 300   | 400   |\n|-----------------------------|-------|-------|-------|-------|-------|\n| **CR$^2$ Avg. Solution Length** | 56.76 | 46.35 | 38.42 | 32.84 | 29.16 |\n| **CR$^2$ Success Rate**     \t| 0.63  | 0.62  | 0.59  | 0.54  | 0.33  |\n| **CRL Avg. Solution Length**\t| 62.96 | 49.88 | 41.94 | 36.11 | 31.77 |\n| **CRL Success Rate**       \t| 0.54  | 0.50  | 0.44  | 0.40  | 0.30  |\n\n**> The \"repetition factor\" hack seems unprincipled - why this specific approach?**\n\nThank you for raising this point. While the repetition factor may appear ad hoc at first glance, it is in fact motivated both empirically and conceptually.\n\nWe tested several natural alternatives for incorporating in-trajectory negatives, and found that our method consistently outperforms them — as shown in Appendix Figures 23 and 24, in both Sokoban and Rubik’s Cube. These comparisons demonstrate that the repetition factor is not just a heuristic convenience, but an empirically validated design choice.\n\nConceptually, the repetition factor ensures that all in-trajectory negatives are properly anchored. In standard contrastive learning (e.g., CRL), negatives (like $y_+$) are simultaneously pulled by their own anchors (e.g., $y$), limiting how far another anchor (e.g., $x$) can push them. Without anchoring, harder in-trajectory negatives (e.g., $x_-$) can drift arbitrarily far due to stronger gradients. Our approach mitigates this by ensuring that all such states appear as anchors in the batch, which stabilizes training and preserves the structure of individual trajectories.\nWe have updated the Method section to clarify this motivation.\n\n**> The conditional independence assumption is crucial to your theoretical analysis. Have you empirically validated this assumption across your test domains? What happens to your theoretical guarantees when this assumption is violated?**\n\nIn the constant context setting, if the environment is fully observable and the context $C$ is deterministically derived from the full state $X$, then the conditional independence assumption holds: observing $X^+$ adds no new information about $C$ beyond $X$. We will clarify this in the revised paper.\nMore generally, the assumption’s validity depends on how $C$ is defined. While it can be trivially satisfied (e.g., $C$ is empty), this is not meaningful in practice.\nIf the assumption is violated, the theoretical guarantees no longer strictly hold; however, our experimental results suggest the method remains practically useful. We have updated the Limitations section to reflect this.\n\n**> Despite the title suggesting general \"combinatorial reasoning,\" testing is limited to puzzle games with no evidence of effectiveness on real combinatorial problems (traveling salesman, vehicle routing, etc.).** \n\nThank you for this suggestion. Our paper focuses on complex, structured reasoning tasks under deterministic dynamics, as in standard benchmarks like the Rubik’s Cube and 15-puzzle, which are widely used to evaluate planning [2, 3, 4] due to their combinatorial complexity.\nWe understand the concern and are open to adjusting the title to better reflect this scope (e.g., Contrastive Representations for Complex Planning Tasks), and would make corresponding edits throughout the paper if preferred.\n\n**> No results across multiple random seeds,**\n\nThank you for pointing this out. Our main results (Figure 4) were averaged over 3 random seeds, though this was not clearly stated in the main text. In some cases, the variance is very low, making the error bars barely visible without zooming in. We have updated the caption of Figure 4 to clarify this.\nTo increase confidence in the results, we are adding more seeds (up to 5) for selected experiments, including those in Figures 4, 8, and 9.\n\n**> insufficient analysis of learned representations beyond basic correlation**\n\nWe thank the reviewer for this comment. To assess whether contextual information is present in the representations, we trained a linear classifier to predict trajectory ID, used as a proxy for static context, from representations produced by CRL, CR$^2$, or a randomly initialized network.\n\nWe used >49k Sokoban transitions, with 70% of each trajectory for training and 30% for testing, and report mean accuracy over 6-fold cross-validation.\n\nResults (mean ± std):\n- CRL: 88.16 ± 0.30%,\n- CR$^2$: 1.80 ± 0.09%,\n- Randomly initialized representations: 3.32 ± 0.15%.\n\nThese results show that CRL retains contextual information, while CR$^2$ effectively suppresses it.\nWe have added a brief summary in Section 4.2 and included full details in the Appendix.\n\n**> missing comparisons to state-of-the-art methods (Graph Neural Networks, neural combinatorial optimization, Monte Carlo Tree Search variants).**\n\nThank you for the suggestion. To the best of our knowledge, state-of-the-art methods such as GNN-based models, neural combinatorial optimization approaches, and MCTS variants have not been applied to the specific puzzle domains we evaluate (e.g., Rubik’s Cube, 15-puzzle, Sokoban) in a way that would allow for direct comparison. If the reviewer is aware of specific methods or references relevant to these domains, we would be grateful for the pointers and would be happy to include them in the paper and, if feasible, in future evaluations.\n\n**> The method makes strong, unrealistic assumptions. It only works if we can cleanly separate \"context\" from \"temporal dynamics\" - a strong assumption that breaks down in many domains. The paper provides no systematic approach for identifying context in practice.**\n\nWe believe there may be some misunderstanding here: Fig. 4 shows experimentally that the method works when the context is not clearly separable. The point of our practical implementation is that it bypasses the need for identifying the context.\nWe agree that the theoretical analysis requires some assumptions (as is common in prior work). We will revise the paper to highlight that these assumptions are not required by our practical method.\n\n**> The related work section is superficial and fails to clearly position the work within existing literature.**\n\nWe have revised the related work section to more clearly highlight the similarities and differences from prior work.\nIf the reviewer could point to specific areas or references we may have overlooked, we would be happy to update the related work section in the paper to better reflect prior work and clarify our contributions.\n\n\n**> The definition of context works well for Sokoban where static elements are clearly separable, but how would you formally define context in domains like Rubik's Cube? Could you provide a more general mathematical framework for identifying context across different combinatorial domains?**\n\nThank you for this insightful question. Our formal definition (Definition 4.1) assumes context corresponds to a clearly separable or static part of the state, as in Sokoban.\n\nIn domains like Rubik’s Cube, where no explicit static structure exists, we treat “pseudo-context” as a latent variable implicitly shared across a trajectory—e.g., abstract features such as the initial configuration or symmetry class. Our empirical approach approximates conditional sampling using in-trajectory negatives.\n\nDeveloping a general mathematical framework for identifying or learning context in such combinatorial domains is an important direction for future work. We will clarify this in the paper.\n\n**> The core insight about negative sample selection is reasonable but offers minimal novelty.**\n\nWe appreciate the reviewer’s recognition that our core insight on negative sample selection is reasonable. While hard negatives are well-established in contrastive learning, our contribution lies in reframing this idea for a setting where contrastive methods typically struggle—combinatorial puzzles.\n\nWe propose a way to leverage trajectory structure to construct informative negatives, enabling success on tasks where prior contrastive methods fail. To our knowledge, this framing and application have not been explored before. If the reviewer is aware of related work, we would be grateful for references so we can properly acknowledge and discuss them.\n\n\n**Summary:**\n\nWe would like to again thank the reviewer for the feedback; we believe that the revised paper will be stronger because of these suggestions and new experiments. We kindly ask that the reviewer reevaluate the paper in light of these revisions. We would be more than happy to run additional experiments, if there are any outstanding concerns about the paper.\n\n**References:**\n\n[1] Kraskov, A., et al., Phys. Rev. E, Estimating Mutual Information\n\n[2] Lehnert, L., et al., COLM 2024, Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping\n\n[3] Bush, T., et al., ICLR 2025, Interpreting Emergent Planning in Model-Free Reinforcement Learning\n\n[4] Zawalski, M., et al., ICLR 2023, Fast and Precise: Adjusting Planning Horizon with Adaptive Subgoal Search"}]

Author rebuttal:

Meta-review: This paper proposes a contrastive learning method that uses in-trajectory negatives to learn representations for combinatorial reasoning tasks, which overcomes a failure mode of standard contrastive approaches.

The reviewers generally agreed that the paper is well-written and nicely motivated the problem by identifying an important limitation in existing methods. They also agreed that the proposed solution is simple yet effective, and the experimental results are convincing. There were are a few initial concerns about the novelty of the proposed method, the limited scope of the evaluation, and a need for more rigorous validation of the theoretical claims about context removal. However, the authors provided additional results from new experiments during the rebuttal period that directly addressed most of the concerns. As a result, a few reviewers increased their scores. 

Given that the the paper got stronger over the rebuttal period, and that the reviewers are unanimously positive about the paper, I recommend acceptance.
Decision: Accept (poster)