refactor(trajectory): iterative leaf DFS and reward assignment moved to the adapter#2174
Open
jingshenghang wants to merge 3 commits into
Open
refactor(trajectory): iterative leaf DFS and reward assignment moved to the adapter#2174jingshenghang wants to merge 3 commits into
jingshenghang wants to merge 3 commits into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Two small, independent refactors to the agent trajectory code, both keeping behavior unchanged.
1.
MessageNode.leaves— recursive → iterative DFSReplace the recursive
yield fromtraversal 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
TrajectoryManagerTrajectoryManageris deliberately tokenizer-free — the adapter decodes.responseafter 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_trajectorydrops itsrewardparameter; every emittedSampleleaves the manager withreward=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.pyandtests/test_agent/test_adapters.pypass (50 passed). Renamedtest_2_8/test_2_9to describe what they actually cover (fork segments / multi-leaf) rather than reward.