Trajectory Graph Learning: Aligning with Long Trajectories in Reinforcement Learning Without Reward Design

We thank the reviewer for recognizing our work's value and providing an excellent summary of our results. Now we address some of the concerns and questions raised by the reviewer.

$**Q1.**$ “Motivation: Why should the agent be closer aligned to the expert trajectory with the cost of removing training data? Some removed trajectories could hold valuable information. For example: Was the tweak with "make the holes passable" in the evaluation necessary to make the algorithm able to solve the task?"

$**A1.**$ Thank you for your question. If we understand correctly, you're asking why our MWIS method selects a subset of expert trajectories—in other words, why we remove some seemingly valuable training data. This is because our objective is to learn a policy that maximizes the probability of matching one of the expert trajectories. While we may have access to a large set of expert demonstrations, many of them may conflict with each other. Intuitively, a single fixed policy cannot simultaneously follow two conflicting trajectories, as discussed in our theoretical analysis. To address this, we employ a graph-based approach to identify a maximum-weight independent set (MWIS), which allows us to select a consistent subset of non-conflicting trajectories for optimizing.

Making the holes passable within our environment was not needed for TGL-UCB algorithm to function, as the TGL algorithm works on the original Frozenlake. The original FrozenLake environment ends an episode when the agent falls into a hole. Making holes passable was just a useful minor modification to allow the failed episodes to continue, which (i) enables longer trajectories for the PPO expert to learn from, (ii) allows for more diverse trajectories and more complex conflicts within our conflict graph structure, and (iii) makes it less trivial to reproduce trajectories. These more diverse and longer (usually full horizon) trajectories are harder to replicate.

$**Q2.**$ “Evaluation: Not reproducable in my opinion. Please explain the experimental setup and evaluation more in detail. For example: what is 'imitation.algorithms.bc'? Is BC and TGL trained on the PPO trajectories?"

$**A2.**$ Thank you for pointing out that the some implementation details were not explicit enough. In our revised appendix and paper, we will add more details such as exact Python packages, seeds, etc., to ensure that our experiment is fully reproducible. Concretely, 'imitation.algorithms.bc' is the behavior-cloning package from the open-source imitation library(https://arxiv.org/abs/2211.11972), which provides the supervised behavior-cloning baseline that we compared against. For the training pipeline, we first train a PPO expert on the modified FrozenLake environment for 50k steps. We roll out this PPO policy to collect a pool of expert demonstration trajectories. TGL and BC both receive this same pool of demonstration trajectories, and no additional information about the environment noise and the PPO expert. The difference is how the two methods (TGL and BC) learn from the demonstration trajectory set. During evaluation, every method is run for 10,000 fresh episodes under the same seed for the environment, and we report the empirical match probability.

$**Q3.**$ “Also the metric is questionable: Does it make sense that the PPO expert is worse in imitating itself? Furthermore: What's the runtime dependency to the problem size? Is it applicable to real-world problems? Maybe the authors could extend the gridsize and measure the runtime. Abstact says "substantially outperforms standard imitation learning methods for long-trajectory planning." but the experiments only include horizons of 10 steps."

$**A3.**$ "Does it make sense that the PPO expert is worse in imitating itself?": At first glance it may seem counter-intuitive that the PPO expert is not the best at reproducing its own demonstrations, but the gap arises because the objectives are different. PPO is trained to maximize its chances of reaching the goal in the noisy FrozenLake environment, not to reproduce any particular action sequence. When we are sampling trajectories from the PPO expert, we ensured that there are no duplicate trajectories, either through environment noise or to sample from the PPO expert stochastically, which can lead to less efficient detours before reaching the goal. BC and TGL are trained explicitly to replicate the demonstration pool behavior, so they are rewarded for copying the detour behavior. Because our evaluation counts only perfect matches (every state and action must be the same as a demonstration trajectory), PPO's performance is limited to the fraction of demonstrations that lie on its preferred set of paths.

Regarding the runtime dependency of our algorithm, we kindly refer the reviewer to Theorems 2 and 4, where the runtime is analyzed in two components: (1) constructing the conflict graph and (2) solving the maximum-weight independent set (MWIS) problem. Both components exhibit polynomial-time complexity with respect to the planning horizon $H$ and the number of expert trajectories $M$. We believe this complexity is reasonable and applicable to certain real-world scenarios.

In response to your suggestion to increase the grid world size, we would like to clarify that the purpose of our experiments is not to demonstrate raw performance in high-dimensional state-action spaces, but rather to isolate and validate the core capability of TGL in solving the trajectory alignment problem using the novel MWIS-based framework. Expanding the grid size to $8 \times 8$ or $10 \times 10$ does increase the number of states, but only by a constant factor. This does not introduce qualitatively new types of trajectory conflicts, nor does it alter the fundamental structure of the conflict graph or the underlying learning dynamics. However, we agree that reporting the runtime could improve transparency, and we are open to including that in future versions for fair comparison.

As for the concern that the planning horizon $H$ is too short, this is primarily due to the limited size of the $4 \times 4$ grid world, which naturally restricts the maximum number of steps. We would also like to clarify a possible misunderstanding regarding our statement in the abstract: when we say our method "substantially outperforms standard imitation learning methods for long-trajectory planning," we are not referring to large values of $H$ per se. Instead, we mean that our method can faithfully imitate entire expert trajectories in a temporally consistent manner—something that standard imitation learning methods often fail to do. The phrase "long-trajectory" here emphasizes exact trajectory alignment over time, rather than the absolute value of $H$.

$**Q4.**$ “All in all I appreciate the mathematically sound problem formulation and method. I would raise my score to accept, if the authors could give a motivation that explains why the algorithm is useful in practice (maybe with examples) and a more elaborated evaluation that leads to more insights and also presents more metrics like the accuracy on the main tasks."

$**A4.**$ Thank you for your appreciation of our theoretical contributions and methodology. Regarding the motivation behind our approach, we offer both analytical and empirical justifications. On the analytical side, we refer the reviewer to Appendix A.2, where we provide concrete examples illustrating how behavior cloning (BC) can fail to capture expert trajectories. The key insight is that traditional methods like BC often rely solely on cumulative rewards and overlook the sequential structure of trajectories. In contrast, our method directly aligns with complete trajectories, thereby capturing richer temporal information. This leads to improved performance on trajectory-level metrics such as matching probability.

On the empirical side, we draw motivation from real-world applications—such as self-driving and large language models (LLMs)—where decision sequences are long, branching, and structurally complex. For example, in self-driving, trajectory planning involves reasoning over distinct, non-overlapping future paths; in LLMs, decoding strategies like beam search produce diverse, tree-structured output sequences. These settings highlight the need to model full trajectories rather than rely solely on scalar feedback.

As for your suggestion regarding more detailed evaluations and metrics, we clarify that our work introduces a new setting—direct trajectory alignment—where the natural evaluation criterion is the probability that a learned policy exactly reproduces the expert trajectory set. Since the goal is replication rather than reward maximization, traditional RL metrics such as episode return or proximity to a ground-truth policy can be uninformative or even misleading. For example, a policy may achieve high returns while never matching any expert trajectories. In our framework, policy quality is judged solely by its alignment with the provided demonstrations, which is why our experiments and metrics are designed around this alignment objective.

$**Q5.**$ “Minor weaknesses"

$**A5.**$ Thank you for pointing out these minor issues, and we will correct them in a revised version.

$**Q6.**$ “Is there a difference between (l143) $d^\pi$ and (l182) $p_i$ for $\tau_i$? Why is it defined two times?"

$**A6.**$ Thank you for your question. In line 143, $d^\pi(\tau)$ denotes the probability of generating trajectory $\tau$ under policy $\pi$—that is, it incorporates both the policy's action distribution and the environment's transition dynamics. In contrast, in line 182, $p_i$ refers to the raw probability of trajectory $\tau_i$ based solely on the product of transition probabilities. These two quantities differ in that $d^\pi(\tau)$ accounts for the stochasticity of the policy, whereas $p_i$ considers only the environment's dynamics.

Thank you for your rebuttal. I will raise my score to Borderline Accept.

We thank the reviewer for recognizing our work's value and providing an excellent summary of our results. Now we address some of the concerns and questions raised by the reviewer.

$**Q1.**$ “The proposed approaches in Section 4 and 5 appear to be well established techniques and their error bounds are also naïve."

$**A1.**$ Thank you for your question. In Section 4, we introduce two practically motivated settings: bounded realizability of the expert trajectory set and tree-structured MDPs. Both settings are grounded in real-world scenarios and, importantly, admit polynomial-time algorithms. While the theoretical bounds we provide may appear relatively simple, the key contribution lies in identifying meaningful and realistic cases in which the otherwise NP-hard MWIS problem becomes tractable. In Section 5, we further develop a novel UCB-style exploration strategy tailored to our TGL framework. To analyze its sample complexity, we reduce the problem to a combinatorial multi-armed bandit (CMAB) setting. This reduction is non-trivial and reflects a careful design choice to make the exploration problem theoretically analyzable while preserving the core structure of trajectory alignment.

$**Q2.**$ “While the paper claims that the method is effective for long-horizon problems, but this central claim lacks sufficient analytical or empirical support. It is not made clear why optimizing the proposed objective function inherently leads to better performance on long-horizon tasks compared to other methods."

$**A2.**$ Thank you for raising the concern regarding the analytical and empirical support for our method’s effectiveness in long-horizon problems. For analytical support, we kindly refer the reviewer to Appendix A.2, where we provide concrete examples illustrating how behavior cloning (BC) fails to capture expert trajectories. The key intuition is that traditional methods, such as BC, often rely solely on cumulative rewards and fail to preserve the sequential structure of trajectories. In contrast, our method directly aligns with entire trajectories, capturing richer and more complete information. This leads to better performance under trajectory-level metrics such as matching probability. On the empirical side, we have also discussed practical motivations in real-world applications—such as self-driving and large language models (LLMs)—where decision sequences are long and branch significantly. For instance, in self-driving, trajectory planning involves reasoning over long, non-overlapping future paths; in LLMs, decoding strategies like beam search generate diverse sequences that resemble tree-structured decisions. These examples further highlight the importance of modeling full trajectories rather than relying on scalar feedback alone.

$**Q3.**$ “The experimental validation is limited. The evaluation is conducted on a single, small-scale grid-world environment and the metric is the trajectory match probability, which is directly related to the method's objective function. This makes it difficult to assess the framework's generalization capabilities or its effectiveness on other standard performance metrics."

$**A3.**$ Thank you for highlighting the scope of our experiment. Our experiment deliberately focuses on a small, discrete grid-world environment because (i) it allows us to control for confounding factors and isolate the effect of our new trajectory-alignment framework and objective, and (ii) the theoretical guarantees we have proven currently assume a finite state-action space, so extending the experiments beyond finite MDPs first requires extending our theoretical work. While the evaluation metric, trajectory match probability, indeed relates to the objective function, this similarity is intentional. In the direct trajectory alignment setting we propose, a policy is only considered "good" when it reproduces the demonstrations, independent of learned or assigned rewards. For example, even if the demonstration trajectory set is objectively "poor" based on rewards, the desired goal of our framework remains the same: match the demonstrations. That said, we agree that supplementing match probability with other metrics will help assess the framework's generalization capabilities. In the revised version of our paper, we will add additional measurements of trajectory quality and conflict-graph density, and error bars as a statistical measure to show the applicability of our framework in larger environments.

$**Q4.**$ “The analysis is limited to deterministic policies and the tabular setting."

$**A4.**$ Thank you for the thoughtful suggestion. First, our paper begins with the deterministic policy setting because it offers a clearer formulation and is more straightforward to analyze. However, extending to stochastic policies is not a fundamental issue in our framework. In fact, for the trajectory alignment problem we consider, the optimal policy can always be found within the set of deterministic policies. One can think of the space of all policies as a polyhedron, where deterministic policies correspond to its vertices. Since the optimization objective is linear with respect to the policy distribution, the optimal solution lies on the boundary of this polyhedron—i.e., at a deterministic policy. Thus, it suffices to focus on deterministic policies.

Second, while our setting assumes a finite action space to ensure polynomial-time tractability, the state space can be infinite. This is because the core computational complexity of the TGL framework arises from learning the trajectory alignment graph, which depends on the number of expert trajectories rather than the total number of states. As a result, in the sample complexity bound of our UCB-style exploration method (Section 5), the number of states does not appear. This shows that our method is not restricted to the tabular setting and can generalize to environments with large or even infinite state spaces.

$**Q5.**$ “In general, the use of direct trajectory alignment is not well motivated. In the example in Appendix A.2, it is unclear to me why the DTA policy is superior than the BC policy. Could you elaborate on when and why DTA is better than BC?"

$**A5.**$ Thank you for your question regarding the example in Appendix A.2. We are happy to provide a more detailed explanation. In this example, we consider two expert trajectories to align with: $\tau_1 = (s_1, a_1, s_2, a_2)$ and $\tau_2 = (s_1, a_2, s_2, a_1)$. Traditional behavior cloning (BC) treats each state-action pair independently. In this case, since both $a_1$ and $a_2$ appear equally often at both $s_1$ and $s_2$ across the two trajectories, BC learns a uniform policy:

where $\pi_1$ and $\pi_2$ are the policies in step 1 and 2 respectively. However, this policy fails to capture the sequential structure of either trajectory. The probability of reproducing each expert trajectory under this policy is:

so the total matching probability for expert trajectories is only:

In contrast, Direct Trajectory Alignment explicitly optimizes the sequence-level probability:

allowing the policy to recover the full structure of either trajectory. For instance, setting $\pi_1(a_1 \mid s_1) = 1$, $\pi_2(a_2 \mid s_2) = 1$ will reproduce $\tau_1$ with probability 1; similarly, $\pi_1(a_2 \mid s_1) = 1$, $\pi_2(a_1 \mid s_2) = 1$ will reproduce $\tau_2$. In both cases, the total expert trajectory

$**Q6.**$ “In numerical experiments, could you present other performance metrics (e.g. distance to a ground-truth expert policy or task performances of the resulting policy). Additionally, how does the policy's performance vary as the horizon length (H) changes?"

$**A6.**$ We thank the reviewer for the suggestion, but would like to clarify that our work introduces a new setting, direct trajectory alignment, for which the natural evaluation criterion is the probability that the resulting policy reproduces the finite expert trajectory set. Because the goal is replication rather than reward maximization, traditional RL metrics such as episode return or distance to an optimal "ground-truth" policy are not informative and could even be misleading. For example, a policy may score highly on those metrics while never matching any demonstration trajectories. In our framework, how good our policy is is determined solely by the quality and diversity of the provided trajectory set, which is why we designed our experiment in this way. As for the second question, by construction, the achievable match probability to the trajectory set does not change when horizon (H) is changed, since the number of trajectories remains fixed. As we change H, the runtime complexity of the resampling operation within our algorithm will scale linearly.

$**Q7.**$ In line 240, I believe the complexity bound should be $M^{1/\epsilon_0}$.

$**A7.**$ Thank you for pointing out the typo. It should be $M^{1/\varepsilon_0}$.

We thank the reviewer for recognizing the value of our work and providing valuable feedback. We address the concerns of the reviewer as follows.

$**Q1.**$ “The paper is relatively theoretical, and some details are not well explained. For instance, the supplementary material includes the full paper with appendices in a single ZIP file, which is not in line with the NeurIPS submission guidelines. Reviewers primarily rely on the officially submitted PDF for their assessment, and any additional technical appendices should ideally be submitted together with the main paper before the full submission deadline to facilitate a more straightforward review process."

$**A1.**$ We thank the reviewer for pointing this out and sincerely apologize for not adhering strictly to the NeurIPS submission guidelines regarding the formatting and inclusion of supplementary materials. We acknowledge that submitting the full paper and appendices as a single ZIP file may have made it more difficult to access key technical details during the review process. We will ensure that all technical appendices are properly integrated into the main supplementary PDF in accordance with the official requirements in any future submissions or revisions. We appreciate the reviewer’s understanding and will take greater care to streamline the presentation for ease of review.

$**Q2.**$ “The theoretical part of the paper is relatively strong, but to enhance the practical significance of the research, it would be valuable to include comparisons with related methods such as Inverse RL and PbRL in the experiments. This would help to demonstrate the practical utility and advantages of the TGL framework beyond its theoretical merits."

$**A2.**$ We appreciate the reviewer’s recognition of the theoretical strength of our paper. Regarding the suggestion to include comparisons with related methods such as Inverse Reinforcement Learning (IRL) and Preference-based Reinforcement Learning (PbRL), we would like to clarify that although these methods are discussed in the motivation and related work sections, our proposed TGL framework addresses a fundamentally different problem—namely, the General Direct Trajectory Alignment problem. Traditional IRL and PbRL approaches typically aim to recover a reward function and evaluate performance through value-based metrics. In contrast, TGL evaluates performance by measuring the matching probability of expert trajectories, which aligns more closely with the objectives of behavior cloning methods. This difference in problem formulation and evaluation criterion is the main reason why we chose to compare TGL against behavior cloning baselines in our experiments.

$**Q3.**$ “Given the current experiments are only conducted on a 4×4 Frozen Lake environment, it would be beneficial to validate the TGL framework on more complex environments such as larger grid-worlds or continuous state/action spaces. For example, environments like the Mujoco robotics simulations or more complex navigation tasks could provide a better understanding of the framework’s generalizability and robustness."

$**A3.**$ We appreciate the reviewer’s suggestion to evaluate our proposed TGL framework in more complex and real-world environments to further assess its scalability and robustness. We would like to clarify that the theoretical foundation of TGL is built on a Markov Decision Process (MDP) with a finite number of states and actions, as well as a finite set of expert trajectories. As such, the current framework does not extend to continuous MDPs or settings involving infinitely many expert trajectories. With respect to larger environments, our experimental objective is not to highlight raw performance in high-dimensional state and action spaces, but rather to isolate and validate TGL’s effectiveness in solving the trajectory alignment problem through our novel MWIS-based method. Scaling the grid to larger sizes—such as $8 \times 8$ or $10 \times 10$—increases the number of states by a constant factor, but does not introduce qualitatively new types of trajectory conflicts. Consequently, the underlying graph structure and learning dynamics of TGL remain unchanged.

$**Q4.**$ “Given the current experiments are only conducted on a 4×4 Frozen Lake environment, it would be beneficial to validate the TGL framework on more complex environments such as larger grid-worlds or continuous state/action spaces. For example, environments like the Mujoco robotics simulations or more complex navigation tasks could provide a better understanding of the framework’s generalizability and robustness."

$**A4.**$ We kindly refer the reviewer to our response to Question 3 (Answer 3), where we have addressed this point in detail.

$**Q5.**$ “Could you include comparisons with related methods such as Inverse RL and PbRL in the experiments? This would help to demonstrate the practical utility and advantages of the TGL framework beyond its theoretical merits."

$**A5.**$ We kindly refer the reviewer to our response to Question 2 (Answer 2), where we have addressed this point in detail.

We thank the reviewer for recognizing the value of our work and providing valuable feedback. We address the concerns of the reviewer as follows.

$**Q1.**$ “The experiments are conducted on relatively simple grid-world tasks. More experiments on larger or continuous MDPs are needed to fully evaluate the scalability and robustness of the proposed method."

$**A1.**$ We appreciate the reviewer’s suggestion to evaluate our proposed TGL framework in more complex and real-world environments to further assess its scalability and robustness. First, we emphasize that the theoretical foundation of TGL is built upon a Markov Decision Process (MDP) with a finite number of states and actions, as well as a finite set of expert trajectories. Therefore, our current framework does not extend to continuous MDPs or settings involving an infinite number of expert trajectories. Regarding larger environments, our experimental objective is not to showcase raw performance in high-dimensional state and action spaces, but to isolate and validate TGL’s capacity to solve the trajectory alignment problem via the novel MWIS-based method. Scaling the grid to larger sizes—such as $8 \times 8$ or $10 \times 10$—increases the number of states by a constant factor but does not introduce qualitatively new types of trajectory conflicts. As a result, the core graph learning structure and learning dynamics of TGL remain consistent.

$**Q2.**$ “The experimental results lack error bars or other statistical significance measures, which makes it difficult to assess the consistency and reliability of the reported performance."

$**A2.**$ We appreciate the reviewer for pointing this out. Our primary metric, the empirical trajectory match probability, is a binomial proportion estimated over 10,000 episodes per evaluation, and each episode is essentially a Bernoulli trial: the generated trajectory either matches or doesn't match a demonstration trajectory. In the worst case: $p = 0.5$, the standard error is $\sqrt{0.5(1-0.5)/10000} = 0.005$, which gives a 95% confidence interval of 0.01. Because of this, we originally decided to omit error bars as the extra numbers might add clutter to our paper. Nevertheless, we recognize that explicit intervals and better explanations make our results more reliable. Our method outperforms BC confidently when the trajectory set is large, and edges past BC when the trajectory set is small, which is expected because of the noise present in the environment (all of the methods are likely close to the maximum probability to match the demonstration trajectories). We will include these details in our revised version of the paper. Additionally, we have revised the Table 1 results in our paper using the empirical match rate as an estimate for p in standard error:

Table 1 – Trajectory-match probability (↑ better) over 10 000 episodes.
Rows indicate the number of deterministic vs. stochastic demos in 𝒟.

$**Q3.**$ “The paper does not appear to have been compiled using the official NeurIPS 2025 LaTeX template, which is a requirement for submission."

$**A3.**$ We sincerely apologize for not using the official NeurIPS 2025 LaTeX template in our initial submission. We appreciate the reviewer’s attention to this requirement. We will promptly recompile the paper using the correct template and ensure that all formatting guidelines are strictly followed in the camera-ready version, should the paper be accepted. Thank you for pointing this out.

$**Q4.**$ “Can the authors conduct experiments on more complex and real-world environments to further validate the scalability and robustness of the proposed TGL framework?"

$**A4.**$ We kindly refer the reviewer to our response to Question 1 (Answer 1), where we have addressed this point in detail.

$**Q5.**$ “Can the authors report error bars or perform statistical significance tests on their experimental results to provide a more rigorous evaluation of the performance of TGL compared to baselines?"

$**A5.**$ We kindly refer the reviewer to our response to Question 2 (Answer 2), where we have addressed this point in detail.

$**Q6.**$ “Can the authors provide more details on the practical relevance and validity of the structural assumptions (bounded realizability of expert trajectories, tree-structured MDP) made in their theoretical analyses?"

$**A6.**$ Thank you for your question. We explain the practical relevance and validity of structural assumptions in our theory here.

A clear example of the bounded realizability assumption can be found in the OpenAI Gym's CartPole or MuJoCo locomotion benchmarks. In these settings, expert trajectories are often generated by a pre-trained high-performing policy (e.g., PPO or SAC). Since these policies reliably solve the task, each recorded expert trajectory corresponds to a successful rollout that occurs with consistent, non-negligible probability. Importantly, because the expert dataset is finite and curated—often by sampling from stable policy rollouts—all trajectories in the set are not only reachable but also share a uniform lower bound on their likelihood under the environment dynamics. This matches the bounded realizability assumption: each trajectory has at least some positive probability of being sampled during interaction, ensuring that the learning agent has a realistic chance of encountering similar trajectories during training.

For tree-structured MDPs, we kindly refer the reviewer to the practical motivation provided in lines 264–269. We reiterate the relevance here for clarity. Tree-structured MDPs naturally arise in scenarios where future decisions evolve independently and paths do not reconverge. For instance, in large language models (LLMs), decoding strategies such as beam search or top-k sampling generate diverse continuations from a given prompt, forming a forward tree where each branch corresponds to a distinct output sequence. Similarly, in autonomous driving, planning algorithms simulate future actions—such as turning or accelerating—under physical and safety constraints that prevent the system from revisiting prior states, yielding a branching structure of possible, non-overlapping trajectories.