Vai al contenuto principale

rStar-Math di Microsoft: guida con implementazione

Scopri come creare un'implementazione semplificata del framework RStar usando una combinazione di reti neurali, ragionamento simbolico e Monte Carlo Tree Search (MCTS).
· 12 min leggi

Scopri con l'IA

Il paper RStar-math di Microsoft presenta un approccio innovativo alla risoluzione di problemi matematici usando una combinazione di reinforcement learning, ragionamento simbolico e Monte Carlo Tree Search (MCTS).

In questo blog esplorerò il framework RStar e i suoi componenti principali. Ti guiderò poi passo dopo passo attraverso un'implementazione semplificata che ne dimostra i concetti chiave usando Gradio. Anche se questa demo è ispirata al paper, alcune complessità sono state semplificate per renderla più accessibile.

Che cos'è rStar-Math di Microsoft?

RStar-math punta a colmare il divario tra ragionamento simbolico e capacità di generalizzazione dei modelli neurali pre-addestrati. Il framework integra componenti come Monte Carlo Tree Search (MCTS), modelli linguistici pre-addestrati e reinforcement learning per consentire un'esplorazione efficiente delle strategie di problem solving.

L'idea centrale è rappresentare il ragionamento matematico come un processo di ricerca su un albero strutturato di possibili passaggi, in cui ogni nodo rappresenta una soluzione parziale o uno stato.

framework rStar math

Fonte: Guan et al., 2025

Alcuni dei motivi che rendono rStar-Math particolarmente interessante per me sono:

  1. Include una rete neurale (modello di policy) che prevede la prossima azione nella risoluzione di un problema matematico, guidando l'esplorazione dell'MCTS.
  2. Una rete (modello di reward) che valuta il successo delle azioni intraprese durante i rollout MCTS e fornisce feedback per l'addestramento.
  3. RStar usa librerie di calcolo simbolico come SymPy per operazioni matematiche precise o ragionamento simbolico, come risolvere equazioni o calcolare derivate.
  4. Integra un algoritmo di Monte Carlo Tree Search che esplora sistematicamente i possibili percorsi di soluzione tramite simulazioni, bilanciando esplorazione (provare nuovi percorsi) ed exploit.
  5. Un meccanismo di feedback in cui i modelli di policy e reward vengono addestrati iterativamente in base agli esiti dei rollout MCTS, migliorando nel tempo il processo decisionale.
  6. Il processo di ragionamento è strutturato come un albero gerarchico in cui i nodi rappresentano stati e gli archi rappresentano transizioni.

Panoramica del progetto demo: risolutore di problemi di matematica con Gradio

La demo mostra come un modello di policy e un modello di reward, combinati con il ragionamento simbolico tramite la libreria sympy, possano affrontare problemi matematici in modo strutturato. Le caratteristiche fondamentali di questa implementazione includono:

  1. Modello di policy: una rete neurale che prevede la prossima azione nel processo di problem solving.
  2. Modello di reward: una rete che valuta il successo delle azioni intraprese durante i rollout MCTS.
  3. Ragionamento simbolico: utilizza SymPy per calcoli matematici precisi e per risolvere equazioni.
  4. Monte Carlo Tree Search: implementa una versione semplificata di MCTS per esplorare efficientemente le soluzioni possibili.
  5. Ciclo di reinforcement learning: un ciclo di addestramento di base per migliorare i modelli di policy e reward in base al feedback.
  6. Supporto per equazioni a una o più variabili: gli utenti possono inserire una o due equazioni per trovare soluzioni per variabili come x e y.

Per mantenere la demo semplice e mirata, alcune funzionalità avanzate menzionate nel paper sono fuori dallo scopo di questo tutorial. Queste funzionalità sono:

  1. Scalabilità: il paper originale usa grandi modelli pre-addestrati e risorse computazionali estese. La demo usa reti neurali più piccole ed evita pre-training complessi.
  2. Strategie MCTS avanzate: tecniche come UCT adattivo e strategie di esplorazione diversificate non sono implementate completamente.
  3. Generalizzazione dei task: l'implementazione si concentra esclusivamente sulla risoluzione di equazioni algebriche, mentre RStar è progettato per generalizzare a compiti matematici più ampi.
  4. Dataset: invece di usare un dataset curato per l'addestramento, la demo si affida al ragionamento simbolico e agli input forniti dagli utenti.

Passo 1: Prerequisiti

La demo è suddivisa in diversi componenti, ognuno dei quali riflette una parte della metodologia RStar. Prima di iniziare, assicurati di avere installato quanto segue:

  • Python 3.8+
  • Librerie richieste: installa i pacchetti Python necessari con pip:
pip install requests gradio, sympy 

Poi importa queste librerie:

import gradio as gr
import numpy as np
import torch
import re
import torch.nn as nn
import torch.optim as optim
from sympy import symbols, Eq, solve, N, sin, cos, tan, exp, log, E, sympify
from random import choice

Ora che tutte le dipendenze sono installate, configuriamo i componenti principali.

Passo 2: Reti neurali per Policy e Reward

Queste reti sono versioni leggere dei modelli descritti nel paper, usate per prevedere la prossima azione e valutare il successo. Il modello di policy prevede i prossimi passaggi per risolvere le equazioni date. Usa una rete neurale feedforward per elaborare rappresentazioni codificate del problema.

Analogamente, il modello di reward valuta soluzioni parziali per guidare il processo MCTS. Entrambi i modelli sono implementati con PyTorch.

class PolicyModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.relu = nn.ReLU()
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x
class RewardModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.relu = nn.ReLU()
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

Successivamente, costruiamo una classe nodo per gli alberi MCTS.

Passo 3: Classe TreeNode per rappresentare gli stati MCTS

La classe TreeNode rappresenta i nodi nell'albero MCTS. Ogni nodo corrisponde a uno stato nel processo di ricerca e contiene:

  • Lo stato (ad es. equazioni o soluzioni parziali).
  • Un riferimento al nodo genitore.
  • Un elenco di nodi figli (stati espansi).
  • Visite e Q-value, che tracciano la frequenza di esplorazione del nodo e le ricompense accumulate.
class TreeNode:
    """Represents a node in the MCTS tree."""
    def __init__(self, state, parent=None):
        self.state = state  # Current state
        self.parent = parent
        self.children = []
        self.visits = 0
        self.q_value = 0.0  # Accumulated rewards
    def is_fully_expanded(self):
        return len(self.children) > 0
    def best_child(self, exploration_weight=1.4):
        """Select the best child using UCT formula."""
        def uct_value(child):
            return (child.q_value / (child.visits + 1e-6)) + exploration_weight * np.sqrt(np.log(self.visits + 1) / (child.visits + 1e-6))
        return max(self.children, key=uct_value)
    def add_child(self, child_state):
        """Add a child node with the given state."""
        child = TreeNode(state=child_state, parent=self)
        self.children.append(child)
        return child

Ora che abbiamo la struttura di base, passiamo ai componenti core della demo.

Passo 4: La classe MathSolver

La classe MathSolver è il cuore della demo e combina ragionamento simbolico con ricerca guidata da reti neurali. Implementa diversi componenti chiave:

A. Modelli di policy e reward

Il PolicyModel prevede i prossimi passaggi per risolvere le equazioni, mentre il RewardModel valuta il successo di soluzioni parziali o complete.

class MathSolver:
    def __init__(self, dataset=None):
        self.dataset = dataset or []  # Dataset of math problems
        self.policy_model = PolicyModel(input_size=128, hidden_size=64, output_size=4)  
        self.reward_model = RewardModel(input_size=128, hidden_size=64, output_size=1)  
        self.policy_optimizer = optim.Adam(self.policy_model.parameters(), lr=0.001)
        self.reward_optimizer = optim.Adam(self.reward_model.parameters(), lr=0.001)
        self.execution_context = {}  

Il metodo sopra inizializza la classe MathSolver configurando i componenti necessari per risolvere problemi matematici. Accetta facoltativamente un dataset di problemi e inizializza due reti neurali: il modello di policy, che prevede la prossima azione, e il modello di reward, che valuta il successo delle azioni. 

Ora abbiamo una funzione di policy e una di reward. Successivamente, dobbiamo analizzare e codificare le equazioni in input.

B. Parsing e codifica delle equazioni

Le equazioni vengono analizzate usando sympy e codificate in vettori di feature per l'elaborazione da parte dei modelli di policy e reward.

def encode_problem(self, problem):
 # Advanced encoding using symbolic representation and problem length
     variables = len(re.findall(r'[a-zA-Z]', problem))
     operators = len(re.findall(r'[\+\-\*/\^]', problem))
      problem_length = len(problem)
      return np.array([variables, operators, problem_length] + [0] * 125)

Il metodo encode_problem converte un problema matematico in una rappresentazione numerica a dimensione fissa per i modelli. Estrae feature come numero di variabili, operatori e lunghezza del problema, codificandole in un array NumPy a 128 dimensioni. Questa rappresentazione cattura la struttura del problema, abilitando un'elaborazione efficace da parte dei modelli.

C. Predizione del modello di policy

Il codice seguente genera i passaggi successivi per risolvere le equazioni fornite, inclusa la definizione delle variabili, la creazione delle equazioni e la loro risoluzione.

def policy_model_predict(self, equation1, equation2=None):
    try:
        equations = []
        if equation1:
            equations.append(sympify(equation1.strip()))  # Sympify only equations
        if equation2:
            equations.append(sympify(equation2.strip()))
        all_variables = set()
        for eq in equations:
            all_variables.update(eq.free_symbols)
        var_definitions = [f"{v} = symbols('{v}')" for v in all_variables]
        steps = [
            ("Define variables", "\n".join(var_definitions)),
            ("Define equation(s)", f"equations = {equations}"),
            ("Solve equation(s)", f"solution = solve(equations, {list(all_variables)})"),
            ("Print solution", "print(solution)")
        ]
        return steps
    except Exception as e:
        print(f"Error during policy model prediction: {e}")
        return []

La funzione policy_model_predict analizza le equazioni in input usando sympify di SymPy per assicurarsi che siano espressioni matematiche valide. Identifica poi tutte le variabili presenti nelle equazioni e le risolve usando la funzione solve di SymPy. Questo metodo funge da guida di alto livello per il flusso di lavoro del problem solving. 

D. Predizione del modello di reward

Il metodo reward_model_predict svolge un ruolo essenziale nel reinforcement learning fornendo feedback per le azioni intraprese durante i rollout del Monte Carlo Tree Search (MCTS). 

  
  def reward_model_predict(self, steps, success):
        encoded_steps = self.encode_problem(str(steps))
        encoded_steps = torch.tensor(encoded_steps, dtype=torch.float32)
        reward = self.reward_model(encoded_steps)
        return reward.item() if success else -reward.item()

La funzione codifica i passaggi di problem solving in formato numerico e li valuta tramite il modello di reward, restituendo una ricompensa positiva in caso di successo e negativa in caso di fallimento. Questo feedback addestra il modello di policy, guidandolo a dare priorità alle azioni efficaci e a migliorare le decisioni. Con le funzioni di predizione di policy e reward pronte, possiamo ora lavorare sull'esecuzione.

E. Funzione di esecuzione del codice

Questo metodo gestisce soluzioni multivariabili come tuple o dizionari e converte i risultati simbolici in approssimazioni numeriche usando la funzione N di SymPy. 

    def execute_code(self, code):
        try:
            # Ensure necessary imports and variables are in the execution context
            exec("from sympy import symbols, Eq, solve, N, sin, cos, tan, exp, log, E", self.execution_context)
            # Dynamically initialize variables in the context
            for var_def in self.execution_context.get('var_definitions', []):
                exec(var_def, self.execution_context)
            exec(code, self.execution_context)
            if "solution" in self.execution_context:
                symbolic_solution = self.execution_context["solution"]
                # Handle multi-variable solutions as tuples
                if isinstance(symbolic_solution, list):
                    self.execution_context["solution"] = [tuple(map(N, sol)) if isinstance(sol, tuple) else N(sol) for sol in symbolic_solution]
                elif isinstance(symbolic_solution, dict):
                    self.execution_context["solution"] = {k: N(v) for k, v in symbolic_solution.items()}
                else:
                    self.execution_context["solution"] = N(symbolic_solution)
            return True
        except Exception as e:
            print(f"Error executing code: {e}")
            return False

Questo metodo garantisce calcoli accurati e supporta una gestione flessibile di vari formati di soluzione. In caso di errori, registra il problema e restituisce False, mantenendo una gestione degli errori robusta.

F. Monte Carlo Tree Search (MCTS)

Il metodo MCTS seleziona iterativamente gli stati migliori, espande l'albero di ricerca e simula possibili soluzioni. Le ricompense delle simulazioni vengono retropropagate per migliorare il processo decisionale.

def mcts(self, equation1, equation2=None, num_rollouts=10):
    root = TreeNode(state=(equation1, equation2))
    for _ in range(num_rollouts):
        # Selection
        node = root
        while node.is_fully_expanded() and node.children:
            node = node.best_child()
        # Expansion
        if not node.is_fully_expanded():
            steps = self.policy_model_predict(*node.state)
            for step, code in steps:
                child_state = (step, code)
                node.add_child(child_state)
        # Simulation
        success = True
        for step, code in steps:
            if not self.execute_code(code):
                success = False
                break
        # Backpropagation
        reward = self.reward_model_predict(steps, success)
        while node:
            node.visits += 1
            node.q_value += reward
            node = node.parent
    return root.best_child().state if root.children else None

Il metodo mcts esegue iterativamente quattro passaggi chiave: 

  • Selezione: aiuta a scendere nell'albero usando i nodi figli migliori.
  • Espansione: in questo passaggio si creano nuovi nodi figli usando il metodo policy_model_predict.
  • Simulazione: vengono eseguite tutte le azioni e in questo passaggio si determina il successo.
  • Backpropagation: le ricompense sono calcolate con il metodo reward_model_predict e propagate verso l'alto nell'albero per aggiornare i valori dei nodi. 

Dopo un numero specificato di rollout, il metodo restituisce lo stato del nodo figlio migliore, che rappresenta la soluzione più promettente esplorata durante la ricerca.

G. Esecuzione della soluzione

Il metodo solve orchestra l'intero processo, dall'analisi delle equazioni all'esecuzione e validazione delle soluzioni.

def solve(self, equation1, equation2=None):
    self.execution_context = {}
    steps = self.policy_model_predict(equation1, equation2)
    variables = set()
    for eq in [equation1, equation2] if equation2 else [equation1]:
        if eq:
            variables.update(sympify(eq.strip()).free_symbols)
    self.execution_context['var_definitions'] = [f"{v} = symbols('{v}')" for v in variables]
    steps_output = ["Best solution found:"]
    for step, code in steps:
        steps_output.append(f"Step: {step}")
        steps_output.append(f"Code: {code}")
        if self.execute_code(code):
            steps_output.append("Execution successful.")
        else:
            steps_output.append("Execution failed.")
    if "solution" in self.execution_context:
        final_answer = self.execution_context["solution"]
        if isinstance(final_answer, dict):
            for var, value in final_answer.items():
                steps_output.append(f"{var} = {value}")
        elif isinstance(final_answer, list):
            for solution in final_answer:
                if isinstance(solution, tuple):
                    for idx, var in enumerate(variables):
                        steps_output.append(f"{list(variables)[idx]} = {solution[idx]}")
                else:
                    steps_output.append(f"Solution: {solution}")
        else:
            steps_output.append(f"Final Answer: {final_answer}")
    else:
        steps_output.append("No final answer found.")
    return "\n".join(steps_output)

Il metodo solve elabora una o due equazioni fornite dall'utente inizializzando un contesto di esecuzione e generando i passaggi tramite policy_model_predict. Esegue ogni passaggio, registra i progressi e riporta esito positivo o negativo. Le soluzioni, incluse quelle a una o più variabili, sono formattate con nomi e valori delle variabili per chiarezza. Se non viene trovata alcuna soluzione, viene mostrato un messaggio appropriato.

Abbiamo tutti i componenti principali, quindi possiamo passare all'app Gradio.

Passo 5: Creare un'interfaccia user-friendly con Gradio

L'interfaccia Gradio consente agli utenti di inserire equazioni (una o più), risolverle e visualizzare i risultati in modo interattivo.

with gr.Blocks() as app:
    gr.Markdown("# Math Problem Solver with Advanced Multi-Step Reasoning and Learning")
    with gr.Row():
        equation1_input = gr.Textbox(label="Enter the first equation (e.g., x + y - 3)", placeholder="x + y - 3")
        equation2_input = gr.Textbox(label="Enter the second equation (optional, e.g., x - y - 1)", placeholder="x - y - 1")
    solve_button = gr.Button("Solve")
    solution_output = gr.Textbox(label="Solution", interactive=False)
    solve_button.click(solve_math_problem, inputs=[equation1_input, equation2_input], outputs=[solution_output])
app.launch(debug=True)

Il codice qui sopra crea un'interfaccia utente Gradio per risolvere equazioni matematiche con ragionamento avanzato. L'interfaccia è racchiusa in un contenitore gr.Blocks, che contiene due campi di input tramite gr.Textbox: uno per la prima equazione (obbligatoria) e un altro per la seconda (opzionale).

L'output viene visualizzato in un singolo gr.Textbox etichettato "Solution". Il comando interface.launch() avvia l'app Gradio in un browser e il flag debug=True abilita log dettagliati per aiutare a risolvere gli errori.

Passo 6: Test e validazione

È il momento di testare la nostra app Math Problem Solver. Ecco alcuni test che ho eseguito:

1. Equazione a singola variabile: ho provato a trovare i possibili valori di una singola variabile x data una sola equazione in input.

Risoluzione di un'equazione con variabile singola

2. Problema con più variabili e più equazioni: ho passato due equazioni con due variabili per trovare i possibili valori delle variabili x e y.

Risoluzione di equazioni con due variabili

Possibili estensioni

 Questa demo è una versione di base di ciò che si può ottenere con le capacità del metodo rStar-math. C'è ancora molto lavoro che si può fare per estenderne le funzionalità. 

  • Aggiungere modelli linguistici pre-addestrati per potenziare le capacità di ragionamento del modello di policy.
  • Integrare strategie MCTS avanzate o provare tecniche di esplorazione diversificate per migliorare efficienza e accuratezza del processo di ricerca.
  • Estendere la demo per risolvere equazioni polinomiali di grado superiore e sistemi di equazioni complessi.
  • Addestrare i modelli su un dataset più ampio di equazioni per una migliore generalizzazione.
  • Estendere la demo per gestire compiti di ragionamento aggiuntivi oltre la matematica.

Puoi fare riferimento al repository originale del paper rStar-math su GitHub.

Conclusione

Questa demo mostra un'implementazione pratica del ragionamento multi-step per la risoluzione di equazioni matematiche. Combinando reti neurali, ragionamento simbolico e MCTS, offre uno sguardo su come le tecniche di IA avanzate possano affrontare compiti di ragionamento strutturato. Miglioramenti futuri potrebbero avvicinarla alle piene capacità del framework RStar.

Aashi Dutt

Autore

Aashi Dutt

Sono una Google Developers Expert in ML (Gen AI), una Kaggle 3x Expert e una Women Techmakers Ambassador con oltre 3 anni di esperienza nel tech. Ho co-fondato una startup health-tech nel 2020 e sto conseguendo un master in informatica al Georgia Tech, con specializzazione in machine learning.

Argomenti

Correlato