Don't Drop the BATON: Long-Horizon Robot Manipulation via Agentic Subtask Exploration and Transition-aware Memory

Overview§

BATON addresses two failure modes in LLM-agent orchestration of frozen VLAs for long-horizon manipulation: (1) whole-task exploration cost grows multiplicatively with stages, and (2) lack of transition representation between subtasks causes silent state incompatibility. BATON makes the subtask the unit of exploration in the cheap short-horizon regime and equips memory with transition-aware mechanisms—verifier, handoff, and lookahead transitions—without updating any policy parameters.

Core Methodology§

Additive Exploration Cost§

Whole-task exploration requires roughly $T^K$ episodes for $K$ stages if each stage needs $T$ trials, and failures are unattributable. BATON explores each subtask independently, storing successful solutions in language memory. Total cost becomes:

$$C_{\text{BATON}} = \sum_{k=1}^{K} T_k \approx T \cdot K$$

Each failure is attributed to a single stage, enabling targeted retries. A long-horizon trajectory is then composed from stored subtask solutions rather than discovered whole.

Transition-Aware Memory§

Within a subtask, a verifier agent gates VLA invocation: the VLA is called only after wrist-view perception confirms the scene is ready. Across subtasks:

• Handoff transition restores an entry state perturbed by the predecessor's residue.
• Lookahead transition selects the strategy whose outcome the successor can inherit.

These transitions are stored as textual memory, enabling the LLM agent to reason about preconditions and postconditions in natural language.

Implementation Sketch§

The agent keeps a language-based memory. For each subtask in the plan, it checks entry conditions (verifier), restores if needed, executes the VLA skill, and validates outcome. Below is a simplified PyTorch-style loop (conceptual):

class BATONAgent:
    def __init__(self, llm, vla, verifier):
        self.llm = llm          # planning and memory reasoning
        self.vla = vla          # frozen contact-rich skill executor
        self.verifier = verifier  # wrist-view readiness check
        self.memory = []

    def execute(self, task):
        plan = self.llm.plan(task)  # list of subtasks
        state = reset_env()
        for subtask in plan:
            # Handoff: restore predecessor residue
            state = self.llm.handoff(state, subtask, self.memory)
            # Verifier: ensure entry condition
            while not self.verifier.check(wrist_view(state), subtask.entry_condition):
                state = self.llm.recover(state, subtask, self.memory)
            # VLA only for contact-rich segment
            state = self.vla.execute(state, subtask.skill)
            # Lookahead: store outcome for successor
            self.memory.append(self.llm.lookahead(state, subtask, self.memory))
        return state

No gradients are computed; all adaptation is in language memory.

Results§

On RoboMemArena, BATON improves task success by 11.6% and cumulative success by 14.9% over the prior state-of-the-art, with no parameter updates—demonstrating that agentic exploration and transition-aware memory are effective alternatives to fine-tuning.

Key Equations§

• Whole-task exploration cost: $\mathcal{O}(T^K)$
• BATON cost: $\mathcal{O}(T \cdot K)$
• Success attribution: each failure maps to a single stage $k$, enabling additive retries.

BATON is a pure agentic/system-level contribution, showing that the bottleneck in long-horizon robot learning is not the policy's per-skill capability but the orchestration and state-transition discipline around it.