fix(rstar): backpropagate rewards from evaluated rollout nodes

Author: liuzhengyang699

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))
+        return asyncio.run(self.solve_async(question))

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))