Trajectory Optimization for AI Agents: Optimizing the Chain, Not Just the Answer

How Robinhood built a production optimization system that uses end-to-end feedback to continuously improve agent planning, tool usage, retrieval quality, and response generation.

Your AI agent's biggest failure isn't a wrong answer, it's the wrong path taken to get there.

For complex user queries, an agent's answer quality depends on a cascading chain of decisions: how the task is decomposed, which tools are called, what evidence is retrieved, and how the final response is synthesized. When any part of that trajectory is weak, the final answer suffers.

At Robinhood, this challenge shows up clearly in our customer service chatbot. The chatbot handles complex customer questions that require multi-step reasoning, policy grounding, and tool-based retrieval. However, because of strict latency requirements, we cannot rely on iterative, ReAct-style loops that spend 15 seconds executing multiple rounds of planning and reflection. The system needs to do most of its work in a single planned pass.

That makes trajectory quality our central optimization target.

This post details the system we built to solve this problem: a Trajectory Optimizer that uses end-to-end, final-answer feedback to systematically improve the planner, tools, and communicator across the full agent chain. Rather than tweaking only the last prompt in the pipeline, treating the symptom, we optimize the sequence of decisions that produced the answer.

In many agent systems, the final answer is only the visible surface area of a much larger process. Under the hood, a successful run depends on a trajectory that typically includes:

1. Decomposition of the user problem

2. Ordering of reasoning steps

3. Tool selection

4. Evidence gathering

5. Answer synthesis

This chain is especially critical in systems where the agent cannot afford to iteratively recover from mistakes. In a single-pass environment, a weak initial plan propagates forward. Missing one key reasoning step leads to missing tool calls. Missing tool calls lead to incomplete evidence. Incomplete evidence leads to a factually weak or operationally incorrect final answer. For example, consider a customer asking, “Why is my account restricted?” A reasonable initial plan might check the restriction details and investigate common causes such as OFAC, pattern day trading (PDT), fraud, incorrect account information, or a bank transfer reversal. But if the plan fails to include B-Notice as a possible restriction type, the agent will never call the document tool to verify whether the required W-9 or related forms have been submitted. The final response may look plausible, but it is built on an incomplete trajectory: a missing planning step led to a missing tool call, which led to an incorrect or incomplete answer.

While it is tempting to optimize only the final response prompt, that often treats the symptom rather than the source. If you want to improve the final answer, you have to improve the chain that generated it.

CX Chatbot uses a single-agent architecture with access to many tools, rather than a slow hierarchy of sub-agents. It solves complex user queries through a three-stage pipeline:

1. Planner Agent: Generates the reasoning plan for how the query should be solved.

2. Tool Layer: Selects and executes tools based on the generated plan.

3. Communicator Agent: Uses the plan and tool outputs to generate the final assistant answer.

Because multi-iteration ReAct is not our default operating mode, we depend entirely on a high-quality plan upfront. To achieve this, the Planner Agent combines two retrieval channels:

• A Knowledge Store: Retrieves task-specific knowledge, such as Robinhood Help Center content and Standard Operating Procedures (SOPs).

• A Plan Example Vectorstore: Retrieves relevant historical plan examples based on the user query to serve as dynamic, few-shot reasoning patterns.

This architecture gave us a highly capable, low-latency planner. But it exposed a broader engineering challenge: once you have a multi-component agent pipeline, how do you improve it systematically over time?

In our setting, the cleanest supervision usually exists only at the end of the run (e.g., factual checks and supervised answer accuracy against golden answers from human experts). But the root cause of a bad answer may sit anywhere in the trajectory.

The Trajectory Optimizer treats these components as optimizable parts of one chain. Instead of asking only, “Was the answer wrong?”, it asks, “Which part of the trajectory should be improved to make future answers better?”

Here are the six design principles that make this system work in production.

While our optimizer supports isolated tuning (e.g., "Communicator-only") and joint tuning, we found that optimizing the Planner first is the highest-leverage intervention.

The Planner defines the structure of the rest of the trajectory. If the plan is incomplete, downstream components are boxed into a lower-quality path. Tool execution and response synthesis can only work with the structure they inherit. Rather than spreading updates evenly across the stack, the optimizer prioritizes the component that shapes everything downstream.

For example, consider a user asking, “How do I use margin to buy more shares?” An incomplete plan may simply retrieve a Help Center article or SOP and generate generic instructions for using margin. But the correct trajectory first needs to determine whether the customer already has margin enabled. If margin is enabled, the agent should provide guidance on how to use available margin to place trades. If margin is not enabled, the agent should instead explain how to apply for margin investing and outline the relevant account requirements. Optimizing only the downstream tool call is not enough: without updating the plan to branch on margin eligibility and enablement status, the agent still lacks the structure needed to produce the right answer.

Trajectory optimization has to go beyond text edits. For each component, the optimizer supports several kinds of improvements:

• Prompt optimization: Rewriting core instructions.

• Example optimization: Modifying, adding, and deleting few-shot examples.

• Vectorstore refinement: Creating and curating the Plan Example Vectorstore.

If retrieval quality is weak, the Planner receives poor demonstrations even if its prompt is flawless. Trajectory optimization must treat prompts, examples, and retrieval assets as a single connected system.

Fine-grained labels are expensive. It is much easier to know if a final answer is correct than to manually label every intermediate reasoning step or tool call.

Our optimizer is designed for this reality. It takes end-to-end signals (like factual failures or mismatches against golden answers) and pushes that signal backward to propose changes to the trajectory-producing components. It solves the "credit assignment" problem without requiring hand-labeled supervision for every intermediate step.

LLM-based optimization becomes prohibitively expensive if every candidate change requires a full offline evaluation run. To keep the loop efficient, we added a gating stage.

When the optimizer proposes an update, we first test it exclusively on the previous iteration’s unsolved cases. Only if the updated agent clears a pass-rate threshold (e.g., >80% solve rate on those previously failed cases) does it move on to broader evaluation. Combined with train/test splits and early stopping, this prevents overfitting and ensures optimization stays tightly coupled to the failures that motivated it.

Failures in agent systems are often systematic. Optimizing based on a single failed case can lead to brittle, overfitted fixes.

To capture broader patterns, the optimizer uses batch optimization with gradient-style methods like ProTeGi (Prompt Optimization with Textual Gradients). By analyzing a batch of failures, it exposes deep trajectory-level failure modes—for instance, realizing the planner consistently skips a validation step for a specific class of user intents, rather than treating it as a one-off anomaly.

In production, regressions matter just as much as gains. An optimizer that only fixes failures will inevitably create "flipped cases", where previously correct answers suddenly become incorrect.

We address this by injecting optimization history into the LLM's context during subsequent tuning rounds. The optimizer is forced to review which changes were attempted, which failure modes improved, and which updates caused regressions in the past. This anti-regression memory helps the optimizer search for changes that solve new cases without destabilizing the rest of the system.

When put together, our automated optimization loop looks like this:

1. Run the current agent pipeline on an evaluation set.

2. Score final answers using end-to-end metrics.

3. Collect failed or weak-performing cases.

4. Generate candidate updates (using textual gradients) for the planner, tools, communicator, or a selected subset.

5. Validate candidates first on the previous iteration’s unsolved cases (Candidate Gating).

6. Advance only candidates that clear the pass-rate threshold.

7. Run broader evaluation (train/test split) with early stopping.

8. Record gains, regressions, and optimization history.

9. Repeat until improvements plateau.

The most important detail? The optimizer is always grounded in observed failures from the actual end-to-end system.

As AI systems move from single-shot text generation to multi-step execution, the central question for engineering teams becomes less about raw model quality and more about how to improve structured behavior over time. In these systems, the answer is merely a byproduct of a trajectory.

Whether that trajectory is explicit (like our Planner-plus-Tools pipeline) or implicit (a tightly coupled agent loop), the challenge is the same: how do you use the downstream feedback you actually have to improve the sequence of reasoning and action?

For CX Chatbot at Robinhood, Trajectory Optimization changed the way we think about agent reliability. By building an optimizer that works backward from end-to-end feedback, prioritizes the planner, efficiently validates candidates, and fiercely guards against regressions, we created a systematic path for scaling agent quality.

• In production agent systems, answer quality is often constrained by the quality of the reasoning trajectory rather than the capabilities of the underlying model.

• Single-pass agent architectures require stronger upfront planning because they cannot rely on multiple rounds of reflection and correction.

• End-to-end answer evaluation can be used to improve intermediate planning, retrieval, tool selection, and communication steps without requiring expensive step-level annotations.

• Retrieval assets, prompts, and few-shot examples should be optimized as a connected system rather than as isolated components.

• Candidate gating, batch optimization, and regression prevention mechanisms help make LLM-driven optimization practical in production environments.

• Treating agent behavior as an optimizable chain creates a systematic framework for improving reliability over time.

Chaofan Wang is a member of Robinhood’s Agentic AI Team, where he works on production AI systems, agent architectures, and optimization frameworks that improve reliability and customer outcomes.

Kirill Dubovitskiy is a member of Robinhood’s Customer Care Team, partnering closely with engineering teams to build and evaluate AI systems that improve customer support operations.

Davide Giovanardi is an engineer focused on agent design, reasoning systems, and scalable approaches to AI-powered customer experiences.

Interested in learning more about building at Robinhood? Check out the careers page.