SOFAI-DT: Return-Conditioned Control for Efficient Verified Reasoning

Reasoning models can improve their accuracy by retrying with feedback, but deciding when another attempt will succeed and when to escalate to a more capable model remains an open control problem. Iterative self-improvement methods refine answers from feedback, yet they typically stop on fixed retry limits that waste computation on stalled trajectories, or on the model's own confidence, which relies on unreliable self-assessment. We study this choice in settings where answers are expensive to produce but cheap to verify, so an external validator can score every candidate. SOFAI-LM operationalizes the retry-versus-escalate decision with a training-free metacognitive controller that pairs a large language model (LLM) with a large reasoning model (LRM), but its controller does not learn from the history of executed attempts. We propose SOFAI-DT, which keeps that pairing, the validator, and the feedback loop while replacing the decision policy with a compact Decision Transformer trained offline on complete trajectories. The policy conditions on a target return, the quality/cost outcome it is asked to hit, together with validator progress, attempt history, and measured costs. Across graph coloring and code debugging under two held-out model pairs, SOFAI-DT raises validator-confirmed success by an unweighted mean of 4.9 percentage points while reducing mean latency by 21.7% relative to SOFAI-LM, and it issues 2.3--2.5 fewer LLM calls per problem. The full policy reaches 77.6% observed success, compared with 74.6% for a two-layer MLP, 75.1% without history, and 73.5% for a tuned heuristic policy.