feat: mcts policy based on trace scheduler by xuangu-fang · Pull Request #1203 · microsoft/RD-Agent

xuangu-fang · 2025-08-25T09:46:31Z

MCTS policy based on trace scheduler

📚 Documentation preview 📚: https://RDAgent--1203.org.readthedocs.build/en/1203/

…ce-scheduler

you-n-g · 2025-09-12T02:49:04Z

                self.trace.set_current_selection(exp.local_selection)
            self.trace.sync_dag_parent_and_hist((exp, prev_out["feedback"]), cur_loop_id)
+            # Notify MCTS scheduler for value backpropagation
+            scheduler = getattr(self.exp_gen, "trace_scheduler", None)


Can we put these logic in exp_gen, and make the framework less complicated.

you-n-g · 2025-09-12T02:50:27Z

        self.trace_scheduler: TraceScheduler = import_class(DS_RD_SETTING.trace_scheduler)(
            DS_RD_SETTING.max_trace_num,
            DS_RD_SETTING.scheduler_temperature,
+            c_puct=DS_RD_SETTING.scheduler_c_puct,


I think we should put it into the specific schedular to make it more general

you-n-g · 2025-09-12T02:54:25Z

+                        scheduler = getattr(self.exp_gen, "trace_scheduler", None)
+                        if isinstance(scheduler, MCTSScheduler):
+                            try:
+                                scheduler.reset()


exp_gen.reset()

you-n-g · 2025-09-12T06:56:57Z

                        logger.log_object(self.trace, tag="trace before restart")
-                        self.trace = DSTrace(scen=self.trace.scen, knowledge_base=self.trace.knowledge_base)
+                        # Reset the trace; MCTS stats will be cleared via registered callback
+                        self.trace.reset()


self.exp_gen.reset()

you-n-g · 2025-09-12T06:57:21Z


        self.uncommitted_experiments: dict[int, DSExperiment] = {}  # loop_id -> DSExperiment

+    def reset(self) -> None:


Trace does not need reset.

you-n-g · 2025-09-12T06:57:50Z

+        if callable(fn):
+            self._on_reset_callbacks.append(fn)
+
+    def reset(self) -> None:


Default do nothing

you-n-g · 2025-09-12T06:58:24Z

Override def reset for ParallelMultiTraceExpGen

you-n-g · 2025-09-12T07:00:06Z


+        # Register MCTS reset callback so trace.reset() can cascade to MCTS stats reset
+        scheduler = getattr(self.exp_gen, "trace_scheduler", None)
+        if isinstance(scheduler, MCTSScheduler):


…ce-scheduler

…rdingly

you-n-g · 2025-10-14T08:27:22Z

+                    reward = 0.0
+
+            # Attribute the reward to the immediate parent leaf that was expanded
+            parent_tuple = trace.dag_parent[new_idx] if 0 <= new_idx < len(trace.dag_parent) else ()


back propogate recursive

find the recursive logic

you-n-g · 2025-10-14T08:39:03Z

+        if trace.sub_trace_count + self.uncommited_rec_status[trace.NEW_ROOT] < self.max_trace_num:
+            return trace.NEW_ROOT
+
+        # Step 2: consider only available leaves (not being expanded)


You can't only calculate UCT on leaves

you-n-g · 2025-10-14T08:44:39Z

+        # Avoid div-by-zero; encourage exploration when visits are small
+        return self.c_puct * prior * math.sqrt(max(1, self.global_visit_count)) / (1 + visits)
+
+    def select(self, trace: DSTrace) -> tuple[int, ...] | None:


handle parallel

handle unfinished running

you-n-g · 2025-10-14T09:16:31Z

+            else:
+                # New root expansion: no parent leaf to credit; optional: keep a pseudo key
+                pass
+        except Exception as e:


you-n-g · 2025-10-14T09:16:53Z

+            for idx in range(start_idx, end_idx):
+                self.observe_feedback(trace, idx)
+            self.last_observed_commit_idx = end_idx
+        except Exception as e:


…method

you-n-g · 2025-10-17T08:22:56Z

+            new_idx: Index of the newly appended experiment in trace.hist.
+            reward: Optional explicit reward. If None, derive from feedback.decision (1.0/0.0).
+        """
+        if reward is None:


remove reward

you-n-g · 2025-10-17T08:23:18Z

+            reward: Optional explicit reward. If None, derive from feedback.decision (1.0/0.0).
+        """
+        if reward is None:
+            if 0 <= new_idx < len(trace.hist):


* init mcts class * full ver of MCTS * auto-lint * make MCTS feedback in exp-gen() * refactor: move reset logic from Trace to ExpGen and update usage accordingly * fix: reinitialize trace on consecutive errors in DataScienceRDLoop * feat: add reset method to BaseScheduler and call in MCTSScheduler reset * style: reorder imports for consistency and PEP8 compliance * lint * fix observe_feedback * fix bug * remove uncommited_rec_status * more simple * refactor: move commit observation logic to process_uncommitted_nodes method * docs: add TODO comment about rule-based virtual root node expansion * add score reward * fix bug * fix small bug * lint * change reward * lint --------- Co-authored-by: Young <afe.young@gmail.com> Co-authored-by: jingyuanlm <842442862@qq.com> Co-authored-by: amstrongzyf <amstrongzyf@126.com>

xuangu-fang added 2 commits August 22, 2025 04:01

init mcts class

2a41471

Merge remote-tracking branch 'origin/main' into fsk/MCTS-based-on-tra…

1dc3616

…ce-scheduler

xuangu-fang changed the title ~~feat: MCTS policy based on trace scheduler~~ feat: mcts policy based on trace scheduler Aug 25, 2025

xuangu-fang added 3 commits September 9, 2025 06:31

Merge remote-tracking branch 'origin/main' into fsk/MCTS-based-on-tra…

f8ec630

…ce-scheduler

full ver of MCTS

2104eca

auto-lint

537d860

you-n-g reviewed Sep 12, 2025

View reviewed changes

make MCTS feedback in exp-gen()

7444df7

you-n-g reviewed Sep 12, 2025

View reviewed changes

you-n-g added 6 commits September 21, 2025 12:21

Merge remote-tracking branch 'origin/main' into fsk/MCTS-based-on-tra…

309d898

…ce-scheduler

refactor: move reset logic from Trace to ExpGen and update usage acco…

cdae43e

…rdingly

fix: reinitialize trace on consecutive errors in DataScienceRDLoop

7f2bfd5

feat: add reset method to BaseScheduler and call in MCTSScheduler reset

22495d1

style: reorder imports for consistency and PEP8 compliance

157d712

lint

8d4ea83

you-n-g reviewed Oct 14, 2025

View reviewed changes

jingyuanlm and others added 9 commits October 14, 2025 15:00

fix observe_feedback

9f8ee5a

fix bug

b53b56a

remove uncommited_rec_status

13e5f49

more simple

953ba5f

refactor: move commit observation logic to process_uncommitted_nodes …

fc2f7dd

…method

docs: add TODO comment about rule-based virtual root node expansion

5475237

add score reward

381930f

fix bug

0e72ef6

fix small bug

18f4eb0

jingyuanlm and others added 3 commits October 17, 2025 03:53

lint

20d6bc0

change reward

b15d5a7

lint

b0195c3

jingyuanlm merged commit 13890e0 into main Oct 17, 2025
9 checks passed

jingyuanlm deleted the fsk/MCTS-based-on-trace-scheduler branch October 17, 2025 08:21

you-n-g mentioned this pull request Oct 17, 2025

chore(main): release 0.8.0 #1030

Merged

you-n-g reviewed Oct 17, 2025

View reviewed changes


		self.uncommitted_experiments: dict[int, DSExperiment] = {} # loop_id -> DSExperiment

		def reset(self) -> None:

Uh oh!

Uh oh!

Conversation

xuangu-fang commented Aug 25, 2025 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

xuangu-fang commented Aug 25, 2025 •

edited by github-actions Bot

Loading