Skip to content

refactor(trajectory): iterative leaf DFS and reward assignment moved to the adapter#2174

Open
jingshenghang wants to merge 3 commits into
THUDM:mainfrom
jingshenghang:jsh/leaves-iterative
Open

refactor(trajectory): iterative leaf DFS and reward assignment moved to the adapter#2174
jingshenghang wants to merge 3 commits into
THUDM:mainfrom
jingshenghang:jsh/leaves-iterative

Conversation

@jingshenghang

@jingshenghang jingshenghang commented Jul 3, 2026

Copy link
Copy Markdown
Collaborator

What

Two small, independent refactors to the agent trajectory code, both keeping behavior unchanged.

1. MessageNode.leaves — recursive → iterative DFS

Replace the recursive yield from traversal with an explicit stack-based DFS, so deep multi-turn trajectories no longer risk hitting Python's recursion limit. Traversal order is preserved (children left-to-right).

2. Reward assignment moved out of TrajectoryManager

TrajectoryManager is deliberately tokenizer-free — the adapter decodes .response after draining a session. This makes it reward-free in the same spirit: the manager only owns what the trajectory looks like, not what each sample scores.

  • get_trajectory drops its reward parameter; every emitted Sample leaves the manager with reward=0.0.
  • finish_session (in the adapter, which already knows the outcome reward) assigns the reward in the same loop that decodes .response.

This also resolves a stale semantic: the manager already assigned the full reward to every sample (not an even split), but the docstring/tests still described splitting. The tests are updated to match and no longer assert per-sample reward, since the manager no longer sets it.

Testing

tests/test_agent/test_trajectory_manager_branching.py and tests/test_agent/test_adapters.py pass (50 passed). Renamed test_2_8/test_2_9 to describe what they actually cover (fork segments / multi-leaf) rather than reward.

Replace recursive yield-from with an explicit stack-based DFS so deep
trajectories do not hit Python's recursion limit.
The manager is deliberately tokenizer-free (the adapter decodes .response
after draining); make it reward-free the same way. get_trajectory drops its
reward param and every emitted Sample now leaves the manager with reward=0.0;
finish_session assigns the reward alongside the existing .response decode.
…rd-free

get_trajectory no longer assigns reward (the adapter does), so the branching
tests no longer verify per-sample reward. Removes the reward wrapper plumbing,
the reward args on get_traj calls, and renames test_2_8/2_9 to describe what
they actually cover (fork segments / multi-leaf), not reward.
@jingshenghang jingshenghang changed the title Jsh/leaves iterative refactor(trajectory): iterative leaf DFS and reward assignment moved to the adapter Jul 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant