diff --git a/optillm/rstar.py b/optillm/rstar.py index 520641b..4b2aa60 100644 --- a/optillm/rstar.py +++ b/optillm/rstar.py @@ -53,7 +53,8 @@ async def expand_async(self, node: Node, action: str) -> Node: logger.debug(f"Expanded node with action: {action}") return child_node - async def simulate_async(self, node: Node) -> float: + async def simulate_async(self, node: Node) -> Tuple[Node, float]: + """Return the evaluated node so the whole rollout can receive its reward.""" current_node = node depth = 0 logger.debug("Starting simulation") @@ -66,7 +67,7 @@ async def simulate_async(self, node: Node) -> float: depth += 1 value = self.evaluate(current_node) logger.debug(f"Simulation complete. Final value: {value}") - return value + return current_node, value async def mcts_async(self, root_state: str) -> List[Node]: root = Node(root_state, None) @@ -83,8 +84,8 @@ async def mcts_rollout_async(self, root: Node): action = random.choice(self.actions) if len(node.children) < len(self.actions): node = await self.expand_async(node, action) - value = await self.simulate_async(node) - self.backpropagate(node, value) + terminal_node, value = await self.simulate_async(node) + self.backpropagate(terminal_node, value) async def solve_async(self, question: str) -> str: self.original_question = question @@ -149,7 +150,8 @@ def expand(self, node: Node, action: str) -> Node: logger.debug(f"Expanded node with action: {action}") return child_node - def simulate(self, node: Node) -> float: + def simulate(self, node: Node) -> Tuple[Node, float]: + """Return the evaluated node so the whole rollout can receive its reward.""" current_node = node depth = 0 logger.debug("Starting simulation") @@ -162,7 +164,7 @@ def simulate(self, node: Node) -> float: depth += 1 value = self.evaluate(current_node) logger.debug(f"Simulation complete. Final value: {value}") - return value + return current_node, value def backpropagate(self, node: Node, value: float): logger.debug("Starting backpropagation") @@ -183,8 +185,8 @@ def mcts(self, root_state: str) -> List[Node]: action = random.choice(self.actions) if len(node.children) < len(self.actions): node = self.expand(node, action) - value = self.simulate(node) - self.backpropagate(node, value) + terminal_node, value = self.simulate(node) + self.backpropagate(terminal_node, value) logger.debug("MCTS complete") return self.extract_trajectories(root) @@ -347,4 +349,4 @@ def solve(self, question: str) -> str: """ Synchronous wrapper for solve_async method. """ - return asyncio.run(self.solve_async(question)) \ No newline at end of file + return asyncio.run(self.solve_async(question)) diff --git a/tests/test_rstar_backpropagation.py b/tests/test_rstar_backpropagation.py new file mode 100644 index 0000000..c038d95 --- /dev/null +++ b/tests/test_rstar_backpropagation.py @@ -0,0 +1,66 @@ +import asyncio +import unittest +from unittest.mock import Mock + +from optillm.rstar import RStar + + +class TestRStarBackpropagation(unittest.TestCase): + def test_sync_rollouts_record_terminal_rewards(self): + for depth in (0, 2): + with self.subTest(depth=depth): + search = RStar("", Mock(), "mock", max_depth=depth, num_rollouts=3) + search.generate_response = Mock(return_value="reasoning step") + search.evaluate = Mock(return_value=0.75) + + trajectories = search.mcts("question") + + self.assertEqual(trajectories[0][0].visits, 3) + self.assertAlmostEqual(trajectories[0][0].value, 2.25) + for trajectory in trajectories: + for node in trajectory: + self.assertGreater(node.visits, 0) + self.assertAlmostEqual(node.value / node.visits, 0.75) + self.assertEqual(search.trajectory_score(trajectory), 0.75) + + +class TestAsyncRStarBackpropagation(unittest.IsolatedAsyncioTestCase): + async def test_parallel_trajectories_keep_their_own_rewards(self): + search = RStar("", Mock(), "mock", max_depth=1, num_rollouts=2) + responses = iter(["1", "2", "3", "4"]) + + async def generate(prompt): + # Both rollouts enter expansion before either completes it. + await asyncio.sleep(0) + return next(responses) + + search.generate_response_async = generate + search.evaluate = Mock(side_effect=lambda node: 10 - float(node.state)) + search.mutual_consistency = Mock(return_value=True) + + trajectories = await search.mcts_async("question") + + self.assertEqual(len(trajectories), 2) + self.assertEqual(trajectories[0][0].visits, 2) + self.assertEqual(trajectories[0][0].value, 13) + for trajectory in trajectories: + for node in trajectory[1:]: + self.assertEqual(node.visits, 1) + self.assertEqual(node.value, 10 - float(trajectory[-1].state)) + best = search.select_final_trajectory(trajectories) + self.assertEqual(best[-1].state, "3") + + async def test_zero_depth_counts_each_rollout_once(self): + search = RStar("", Mock(), "mock", max_depth=0, num_rollouts=3) + + async def generate(prompt): + await asyncio.sleep(0) + return "reasoning step" + + search.generate_response_async = generate + search.evaluate = Mock(return_value=0.5) + trajectories = await search.mcts_async("question") + + self.assertEqual(trajectories[0][0].visits, 3) + self.assertEqual(trajectories[0][0].value, 1.5) + self.assertTrue(all(search.trajectory_score(t) == 0.5 for t in trajectories))