Reproduce “REX: Rapid Exploration and eXploitation for AI Agents” on OpenResearch

About this paper

Rithesh Murthy, Shelby Heinecke, Juan Carlos Niebles, Zhiwei Liu, Le Xue, Weiran Yao, Yihao Feng, Zeyuan Chen and 7 more

Salesforce Research developed REX (Rapid Exploration and eXploitation), a framework that integrates an accelerated Monte Carlo Tree Search with Large Language Models to improve AI agent decision-making. REX-UCB achieved 80% accuracy on 2-step Blocksworld tasks, significantly outperforming CoT (40%), while REX-UCL reached 90.44% accuracy on GSM8K-test, surpassing Reflexion (88.85%), all with reduced computational queries compared to traditional search methods.

Key insights

• Accelerating Monte Carlo Tree Search by having an LLM generate an entire solution sequence in a single pass, then backpropagating rewards, drastically reduces computational complexity compared to traditional step-by-step MCTS.
• Explicitly integrating UCB-based feedback into the LLM's prompt (REX-UCB) or directly modifying logits (REX-UCL) systematically guides exploration and exploitation, enabling agents to discover more effective problem-solving pathways.
• The REX-UCB and REX-R variants offer broad compatibility with LLM APIs, even those without access to internal logit values, by leveraging in-context learning to provide reward signals.

Summary by alphaXiv.