Thor Tool Hierarchical Optimization Rl

Outcome: Train Tool-Using LLMs for Robust Mathematical Reasoning

THOR achieves state-of-the-art mathematical reasoning by combining reinforcement learning optimization at two levels: episode-level (correct final answer) and step-level (successful tool execution). This hierarchical approach solves the sparse reward problem inherent in long reasoning chains by recognizing that intermediate tool call success strongly predicts final answer correctness.

Problem Context

Large language models struggle with high-precision mathematical tasks requiring numerical computation and formal symbolic manipulation. Traditional approaches train on supervised examples, missing opportunities to optimize the specific behavior patterns that lead to correct answers. Reinforcement learning offers promise but faces two critical challenges:

1. Sparse rewards: In multi-step reasoning, only the final answer determines success, leaving most intermediate steps without signal
2. Data-policy mismatch: Training on trajectories generated by other models leads to distribution shift and performance degradation

THOR addresses these through a two-phase framework: first constructing aligned training data via TIRGen, then applying hierarchical RL optimization during training and inference.

Core Concept

The key insight is that intermediate tool call success is a strong predictor of final answer correctness. Rather than optimizing only at the episode level (answer correctness), THOR introduces step-level optimization that specifically improves code generation quality at execution failure points. This provides fine-grained reward signal even in long reasoning chains.

During inference, a self-correction mechanism leverages immediate tool feedback: when code execution fails, the model backtracks and explores alternative reasoning paths instead of proceeding with erroneous steps.

Architecture Overview

TIRGen Data Construction Pipeline

• Generator-refiner framework with two agents
• Generator produces reasoning steps (thoughts and tool calls)
• Refiner identifies executable operations and converts them to runnable code
• Ensures data remains in-distribution, preventing policy-data mismatch
• Produces training trajectories aligned with model capability

Hierarchical RL Strategy

• Episode-level optimization: Maximizes final answer correctness using Group Relative Policy Optimization (GRPO)
• Step-level optimization: Applies fine-grained feedback to steps where code generation failed
• Joint optimization of both levels during training
• Trajectory filtering removes execution failures to stabilize gradients
• Backtracking procedure increases action diversity during error correction

Self-Correction Inference Mechanism

• Real-time monitoring of tool execution feedback
• Backtracks when code execution fails
• Regenerates alternative reasoning paths rather than propagating errors
• Operates without requiring retraining or additional inference overhead

Implementation

Phase 1: Data Construction with TIRGen

The TIRGen pipeline creates training data through an iterative generator-refiner loop that ensures tool calls are actually executable.

# TIRGen Pipeline: Multi-agent data construction
import json
from typing import List, Dict, Tuple

class TIRGenPipeline:
    """Constructs tool-integrated reasoning datasets with generator-refiner loop."""

    def __init__(self, generator_model, refiner_model, tool_registry):
        self.generator = generator_model
        self.refiner = refiner_model
        self.tools = tool_registry

    def generate_reasoning_trajectory(self, problem: str) -> Dict:
        """
        Step 1: Generator produces reasoning steps and tool calls.
        Returns trajectory with thoughts and action specifications.
        """
        trajectory = {
            'problem': problem,
            'steps': [],
            'actions': [],
            'observations': []
        }

        # Generator creates reasoning path
        generator_prompt = f"""Solve this problem step by step.
        Problem: {problem}

        Format each step as:
        Thought: [reasoning]
        Action: [tool_name(arguments)]
        """

        response = self.generator.generate(generator_prompt, max_tokens=2048)

        # Parse steps and actions
        for line in response.split('\n'):
            if line.startswith('Thought:'):
                trajectory['steps'].append(line[8:].strip())
            elif line.startswith('Action:'):
                trajectory['actions'].append(line[7:].strip())

        return trajectory

    def refine_trajectory(self, trajectory: Dict) -> Dict:
        """
        Step 2: Refiner converts tool calls to executable code.
        Validates executability and maintains in-distribution samples.
        """
        refined_actions = []

        for action in trajectory['actions']:
            # Extract tool call: tool_name(arg1=val1, arg2=val2)
            tool_name, args = self._parse_action(action)

            if tool_name not in self.tools:
                continue

            # Refiner generates executable code
            refiner_prompt = f"""Convert to executable Python:
            Tool: {tool_name}
            Arguments: {args}
            Available tools: {list(self.tools.keys())}
            """

            code = self.refiner.generate(refiner_prompt, max_tokens=256)

            # Validate executability by dry-run
            if self._validate_code(code, tool_name):
                refined_actions.append({
                    'tool': tool_name,
                    'code': code,
                    'args': args
                })

        trajectory['refined_actions'] = refined_actions
        return trajectory

    def _parse_action(self, action_str: str) -> Tuple[str, Dict]:
        """Extract tool name and arguments from action string."""
        # action_str: "calculator(expression='2+2')"
        tool_name = action_str.split('(')[0]
        args_str = action_str.split('(')[1].rstrip(')')
        args = {}
        for pair in args_str.split(','):
            if '=' in pair:
                k, v = pair.split('=')
                args[k.strip()] = v.strip().strip("'\"")
        return tool_name, args

    def _validate_code(self, code: str, tool_name: str) -> bool:
        """Check if code is syntactically valid and uses correct tool."""
        try:
            compile(code, '<string>', 'exec')
            return tool_name in code
        except SyntaxError:
            return False

Phase 2: Hierarchical RL Training

Episode-level optimization maximizes final answer correctness. Step-level optimization fixes execution failures through targeted policy adjustments.

# Hierarchical RL Training: Episode and Step-Level Optimization
import torch
from torch.optim import AdamW

class HierarchicalRLTrainer:
    """Combines episode-level and step-level RL optimization."""

    def __init__(self, model, tool_executor, learning_rate=1e-5):
        self.model = model
        self.executor = tool_executor
        self.optimizer = AdamW(model.parameters(), lr=learning_rate)

    def compute_episode_reward(self, trajectory: Dict, ground_truth: str) -> float:
        """
        Episode-level reward: 1.0 if final answer matches ground truth, 0.0 otherwise.
        This is sparse but fundamental to solution quality.
        """
        final_answer = trajectory.get('final_answer', '')
        return 1.0 if final_answer.strip() == ground_truth.strip() else 0.0

    def compute_step_rewards(self, trajectory: Dict) -> List[float]:
        """
        Step-level rewards: Track tool execution success at each step.
        Success of intermediate tool calls predicts final correctness.
        """
        step_rewards = []

        for action in trajectory.get('refined_actions', []):
            code = action['code']
            tool = action['tool']

            # Execute code and check for runtime errors
            try:
                result = self.executor.execute(code)
                # Execution success = 1.0
                step_rewards.append(1.0)
            except (RuntimeError, ValueError, KeyError):
                # Execution failure = 0.0 (target for improvement)
                step_rewards.append(0.0)

        return step_rewards

    def compute_grpo_loss(self, batch_trajectories: List[Dict],
                         batch_answers: List[str]) -> torch.Tensor:
        """
        Group Relative Policy Optimization (GRPO) for episode-level training.
        Computes relative rewards within a batch to stabilize gradients.
        """
        batch_size = len(batch_trajectories)
        episode_rewards = [self.compute_episode_reward(traj, ans)
                          for traj, ans in zip(batch_trajectories, batch_answers)]

        # Compute group-relative rewards (normalize within batch)
        episode_rewards_tensor = torch.tensor(episode_rewards, dtype=torch.float32)
        mean_reward = episode_rewards_tensor.mean()
        std_reward = episode_rewards_tensor.std() + 1e-8
        normalized_rewards = (episode_rewards_tensor - mean_reward) / std_reward

        # Generate log probabilities for each trajectory
        log_probs = []
        for traj in batch_trajectories:
            # Reconstruct full token sequence from trajectory
            tokens = self._trajectory_to_tokens(traj)
            log_prob = self.model.compute_log_prob(tokens)
            log_probs.append(log_prob)

        log_probs_tensor = torch.stack(log_probs)

        # GRPO loss: negative of reward-weighted log probability
        loss = -(normalized_rewards * log_probs_tensor).mean()
        return loss

    def compute_step_level_loss(self, batch_trajectories: List[Dict]) -> torch.Tensor:
        """
        Step-level optimization: Fine-grained correction of failed code generation.
        Targets steps where execution failed and uses backtracking to increase diversity.
        """
        step_losses = []

        for trajectory in batch_trajectories:
            step_rewards = self.compute_step_rewards(trajectory)

            # Find failure points (reward = 0.0)
            for i, reward in enumerate(step_rewards):
                if reward == 0.0:  # Execution failed at this step
                    # Extract the generated code and ground truth correction
                    failed_action = trajectory['refined_actions'][i]
                    failed_code = failed_action['code']

                    # Generate alternative code through backtracking
                    # (model explores different code for same tool call)
                    alternative_codes = self._generate_alternatives(
                        trajectory, i, num_alternatives=3
                    )

                    # Compute loss pushing away from failed code
                    failed_log_prob = self.model.compute_log_prob(failed_code)

                    # Encourage alternatives that execute successfully
                    for alt_code in alternative_codes:
                        if self._is_executable(alt_code):
                            alt_log_prob = self.model.compute_log_prob(alt_code)
                            step_losses.append(failed_log_prob - alt_log_prob)

        if step_losses:
            return torch.stack(step_losses).mean()
        else:
            return torch.tensor(0.0)

    def training_step(self, batch_trajectories: List[Dict],
                     batch_answers: List[str]) -> float:
        """
        Combined training step: episode-level + step-level optimization.
        Alternating focus increases both answer correctness and execution robustness.
        """
        # Compute losses
        episode_loss = self.compute_grpo_loss(batch_trajectories, batch_answers)
        step_loss = self.compute_step_level_loss(batch_trajectories)

        # Joint optimization with tuned weights
        total_loss = episode_loss + 0.5 * step_loss

        # Gradient update
        self.optimizer.zero_grad()
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
        self.optimizer.step()

        return total_loss.item()

    def _trajectory_to_tokens(self, trajectory: Dict) -> torch.Tensor:
        """Convert trajectory (thoughts, actions, observations) to token IDs."""
        text = ""
        for step, action, obs in zip(
            trajectory['steps'],
            trajectory['refined_actions'],
            trajectory.get('observations', [])
        ):
            text += f"Thought: {step}\nAction: {action['code']}\nObservation: {obs}\n"
        return self.model.tokenize(text)

    def _generate_alternatives(self, trajectory: Dict, step_idx: int,
                               num_alternatives: int = 3) -> List[str]:
        """Generate alternative code implementations for a failed step."""
        context = "\n".join([str(s) for s in trajectory['steps'][:step_idx+1]])
        tool_name = trajectory['refined_actions'][step_idx]['tool']

        alternatives = []
        for _ in range(num_alternatives):
            prompt = f"""Generate alternative executable code for tool call:
            Context: {context}
            Tool: {tool_name}

            Important: Return only the code, no explanation."""

            code = self.model.generate(prompt, max_tokens=256)
            if self._is_executable(code):
                alternatives.append(code)

        return alternatives

    def _is_executable(self, code: str) -> bool:
        """Check if code is syntactically valid."""
        try:
            compile(code, '<string>', 'exec')
            return True
        except SyntaxError:
            return False

Phase 3: Self-Correcting Inference

During inference, the model monitors tool execution and backtracks when errors occur, enabling dynamic error recovery without retraining.

# Self-Correcting Inference: Real-time Backtracking and Error Recovery
from collections import deque

class SelfCorrectingAgent:
    """Inference-time agent with immediate feedback and dynamic backtracking."""

    def __init__(self, model, tool_executor, max_backtrack_depth=3):
        self.model = model
        self.executor = tool_executor
        self.max_backtrack_depth = max_backtrack_depth

    def solve(self, problem: str, max_steps: int = 20) -> str:
        """
        Solve problem with self-correction.
        If tool execution fails, backtrack and explore alternatives.
        """
        # Stack to track reasoning paths for backtracking
        reasoning_stack = deque(maxlen=self.max_backtrack_depth)

        state = {
            'problem': problem,
            'thoughts': [],
            'actions': [],
            'observations': [],
            'failures': 0
        }

        step = 0
        while step < max_steps:
            # Step 1: Generate next thought and action
            thought, action_code = self._generate_step(state)

            if not action_code:  # Model decided to output answer
                return self._extract_answer(state)

            state['thoughts'].append(thought)
            state['actions'].append(action_code)

            # Step 2: Execute action with immediate feedback
            try:
                observation = self.executor.execute(action_code)
                state['observations'].append(observation)
                step += 1

                # Successfully executed: save state for potential backtracking
                reasoning_stack.append(state.copy())

            except Exception as e:
                # Execution failed: attempt recovery
                state['failures'] += 1
                error_msg = str(e)

                if state['failures'] <= self.max_backtrack_depth:
                    # Backtrack: remove last failed step
                    state['thoughts'].pop()
                    state['actions'].pop()

                    # Re-generate with error feedback
                    error_context = f"Previous attempt failed: {error_msg}\nTry a different approach."
                    state['observations'].append(error_context)

                    # Next iteration will generate alternative code
                    step += 1
                else:
                    # Too many failures: return best current answer
                    return self._extract_answer(state)

        # Reached max steps
        return self._extract_answer(state)

    def _generate_step(self, state: Dict) -> Tuple[str, str]:
        """
        Generate next thought and action code.
        Returns (thought, action_code) or (thought, None) to signal done.
        """
        # Build prompt from current state
        history = ""
        for t, a, o in zip(state['thoughts'][-5:],
                          state['actions'][-5:],
                          state['observations'][-5:]):
            history += f"Thought: {t}\nAction: {a}\nObservation: {o}\n"

        prompt = f"""Problem: {state['problem']}

{history}

Next:
Thought: [reasoning for next step]
Action: [executable Python code, or "ANSWER: value" to output result]
"""

        response = self.model.generate(prompt, max_tokens=512)

        # Parse response
        lines = response.strip().split('\n')
        thought = ""
        action = ""

        for line in lines:
            if line.startswith('Thought:'):
                thought = line[8:].strip()
            elif line.startswith('Action:'):
                action = line[7:].strip()

        # Check if model output answer
        if action.startswith("ANSWER:"):
            return thought, None

        return thought, action

    def _extract_answer(self, state: Dict) -> str:
        """Extract final answer from state observations."""
        for obs in reversed(state['observations']):
            if isinstance(obs, (int, float, str)):
                return str(obs)
        return "No solution found"

Practical Guidance

Hyperparameters and Configuration

When to Use THOR

Use THOR when:

• Training LLMs to solve mathematical or symbolic reasoning tasks
• Problems require precise tool calls (calculators, symbolic solvers, code interpreters)
• Dataset is limited and requires efficient optimization signal (step-level rewards reduce sparse reward problem)
• Model exhibits execution errors that could be corrected dynamically (self-correction provides immediate benefit)
• Multi-step reasoning chains are common (hierarchical approach directly addresses this pattern)
• You have computational budget for RL training (more involved than supervised fine-tuning)

When NOT to Use THOR

• Language generation tasks without tool integration: THOR's benefits depend on tight coupling with tool feedback. For pure language tasks (translation, summarization), supervised fine-tuning is simpler.
• Single-step or fully deterministic problems: If reasoning doesn't branch or fail, step-level optimization provides minimal benefit.
• Extremely large models (>100B parameters): RL training overhead becomes prohibitive; consider simpler policy gradient methods or behavioral cloning.
• No access to execution feedback: THOR requires real-time tool execution results. Without immediate rewards, episode-level RL alone is preferable.
• Real-time inference critical: Self-correction mechanism adds latency (multiple generation attempts per step). For latency-sensitive deployments, use standard inference.
• Data distribution already clean: If TIRGen-quality data is already available, simpler supervised training may suffice without RL overhead.

Common Pitfalls and How to Avoid Them

1. Training on failed trajectories: Include step-level filtering to remove execution failures before gradient updates. Failed steps create misaligned gradients. THOR handles this explicitly; ensure your implementation excludes broken intermediate steps.

2. Ignoring the step-level signal: Episode-level rewards alone recreate the sparse reward problem. Always compute step-level rewards from tool execution failures; set step-level weight ≥ 0.3.

3. Backtracking without diversity: When correcting failed code, ensure the model regenerates with explicit error context. Without this signal, backtracking loops repeating the same failure. Include error messages in the prompt for alternative generation.

4. Generator-refiner data pipeline skipped: Manually created reasoning data often contains non-executable tool calls. Use TIRGen's refiner component to validate executability. In-distribution data is critical for policy alignment.

5. Over-filtering trajectories: Filtering too aggressively (removing all partial failures) eliminates valuable learning signal. Keep trajectories with execution failures at the step level; use them for step-level optimization.

6. Reward signal collision: Episode reward (binary: 0/1) may not distinguish between "almost correct" and "completely wrong." Consider adding intermediate rewards (e.g., partial credit for getting 80% of numerical answer correct) to enrich signal.

7. Batch size too small for GRPO: Group Relative Policy Optimization requires sufficient batch diversity for robust normalization. Use batch size ≥ 32; smaller batches risk noisy gradient estimates.

Reference

THOR is published at ICLR 2026. For implementation details and code, refer to the official repository and paper:

• Paper: https://arxiv.org/abs/2509.13761
• Official Implementation: https://github.com/JingMog/THOR
• Key Technique: Group Relative Policy Optimization (GRPO) extends standard policy gradient with batch-level reward normalization
• Related Work: Chain-of-Thought (Wei et al.), In-Context RL (Rafailov et al.), Tool-Integrated Reasoning (Yao et al. ReAct)