agent_bilevel explorer: recover tool-validated capture, seed info-gain search

When the agent submits + simulator-validates a goal-reaching plan via
evaluate_option_plan / refine_plan_sketch but ends with a prose summary
whose final text doesn't parse into a sketch, the explorer previously
dropped that productive solve to the random-options fallback (a wasted
exploration attempt that could block train-driven early stopping). Enable
the capture path (keyed to the explore task) and, on an empty parse,
rebuild the sketch from the captured plan, grafting each option's
continuous params onto SketchStep.initial_params so they seed the
info-gain parameter search as candidates rather than being replayed
verbatim. Mirrors the test solver's preference for the tool-validated
capture over the final text; stale capture is cleared at entry so an
exploration plan can't leak into a later test solve.

Author: yichao-liang

predicators/explorers/agent_bilevel_explorer.py

@@ -15,6 +15,7 @@
 import logging
 from typing import Any, Callable, Dict, List, Optional, Set
 
+import numpy as np
 from gym.spaces import Box
 
 from predicators import utils
@@ -71,12 +72,20 @@ def _get_exploration_strategy(self, train_task_idx: int,
         # task. Without this, the agent tunes/validates its exploration plan
         # against the wrong task (e.g. a test goal referencing objects this
         # task doesn't even have), so parameter search is meaningless and
-        # only tasks solvable without tuning get solved. The explorer does
-        # not use the approach's capture path, so disable it (and clear any
-        # stale capture) to keep an exploration plan from leaking into the
-        # next test solve.
+        # only tasks solvable without tuning get solved.
+        #
+        # Enable the capture path too (keyed to current_task == this explore
+        # task): the agent often submits + simulator-validates a goal-reaching
+        # plan via evaluate_option_plan / refine_plan_sketch but then ends with
+        # a prose summary whose final text doesn't parse into a sketch. Without
+        # capture that productive solve is lost to the random-options fallback;
+        # with it we recover the captured plan below (see _sketch_from_capture)
+        # and feed its continuous params into the info-gain search. Clear any
+        # stale capture first; the next test _solve re-points current_task and
+        # clears capture again, so an exploration plan can't leak into a test
+        # solve.
         self._tool_context.current_task = task
-        self._tool_context.capture_goal_reaching_plans = False
+        self._tool_context.capture_goal_reaching_plans = True
         self._tool_context.solved_plan = None
         self._tool_context.solved_sketch = None
 
@@ -109,6 +118,16 @@ def _get_exploration_strategy(self, train_task_idx: int,
                 parse_continuous_params=CFG.
                 agent_bilevel_use_llm_initial_params,
             )
+            if not sketch:
+                # The agent's final message didn't parse into a sketch, but it
+                # may have submitted + simulator-validated a goal-reaching plan
+                # via evaluate_option_plan / refine_plan_sketch (captured into
+                # solved_plan / solved_sketch). Recover that plan as the
+                # sketch, carrying its continuous params as initial_params so
+                # they seed the info-gain search below rather than being
+                # replayed verbatim. This mirrors the test solver's preference
+                # for the tool-validated capture over the final text.
+                sketch = self._sketch_from_capture() or []
             if not sketch:
                 raise ValueError("parsed empty plan sketch")
 
@@ -244,6 +263,45 @@ def _get_exploration_strategy(self, train_task_idx: int,
     # Helpers
     # ------------------------------------------------------------------ #
 
+    def _sketch_from_capture(
+            self) -> Optional[List[bilevel_sketch.SketchStep]]:
+        """Rebuild a sketch from a captured, tool-validated plan, or None.
+
+        ``evaluate_option_plan`` / ``refine_plan_sketch`` stash a
+        refined + forward-validated, goal-reaching plan on the explore
+        task into ``solved_plan`` (grounded options with continuous
+        params) and ``solved_sketch`` (the option skeleton plus the
+        subgoals that actually held). We reconstruct a sketch from that
+        skeleton and graft each captured option's continuous params onto
+        the step's ``initial_params``, so the info-gain refinement below
+        seeds them as the first candidate in each step's pool (see
+        ``_sample_info_seeking``) instead of replaying them verbatim.
+        Consume (clear) the capture so it can't be reused.
+        """
+        plan = self._tool_context.solved_plan
+        captured_sketch = self._tool_context.solved_sketch
+        self._tool_context.solved_plan = None
+        self._tool_context.solved_sketch = None
+        if not plan or not captured_sketch:
+            return None
+        seeded: List[bilevel_sketch.SketchStep] = []
+        for i, step in enumerate(captured_sketch):
+            params = None
+            if i < len(plan):
+                params = np.asarray(plan[i].params, dtype=np.float32)
+            seeded.append(
+                bilevel_sketch.SketchStep(
+                    option=step.option,
+                    objects=step.objects,
+                    subgoal_atoms=step.subgoal_atoms,
+                    subgoal_neg_atoms=step.subgoal_neg_atoms,
+                    initial_params=params))
+        logging.info(
+            "agent_bilevel explorer: final text didn't parse, recovered the "
+            "agent's tool-validated plan from capture (%d steps); seeding its "
+            "continuous params into the info-gain search.", len(seeded))
+        return seeded
+
     def _wrap_policy(
             self, policy: Callable[[State],
                                    Action]) -> Callable[[State], Action]:

tests/explorers/test_agent_bilevel_explorer.py

@@ -8,6 +8,7 @@
 from gym.spaces import Box
 
 from predicators import utils
+from predicators.agent_sdk.bilevel_sketch import SketchStep
 from predicators.agent_sdk.tools import ToolContext
 from predicators.explorers import create_explorer
 from predicators.explorers.agent_bilevel_explorer import AgentBilevelExplorer
@@ -302,6 +303,115 @@ def test_plan_truncates_at_deepest_subgoal_failure_after_backtracking():
         f"giving up, got {executed_names}")
 
 
+def _make_captured(pick_params, place_params):
+    """Build the (solved_plan, solved_sketch) a tool capture would stash."""
+    grounded_plan = [
+        _Pick.ground([_block0], np.array(pick_params, dtype=np.float32)),
+        _Place.ground([_block0, _block1],
+                      np.array(place_params, dtype=np.float32)),
+    ]
+    captured_sketch = [
+        SketchStep(option=_Pick, objects=[_block0], subgoal_atoms=None),
+        SketchStep(option=_Place,
+                   objects=[_block0, _block1],
+                   subgoal_atoms={GroundAtom(_On, [_block0, _block1])}),
+    ]
+    return grounded_plan, captured_sketch
+
+
+def test_recovers_captured_plan_when_final_text_unparseable():
+    """Agent validates a plan via evaluate_option_plan but ends in prose:
+
+    explorer recovers the captured plan instead of falling back to
+    random, and seeds the captured continuous params into refinement.
+    """
+    _reset_config()
+
+    goal_state = _make_state({_block0: [0.5, 0.6, 0.0]})
+    option_model = MagicMock()
+    option_model.get_next_state_and_num_actions.return_value = (goal_state, 3)
+
+    pick_params, place_params = [0.42], [0.11, 0.22]
+    grounded_plan, captured_sketch = _make_captured(pick_params, place_params)
+
+    explorer, tool_context = _make_explorer(option_model, None)
+
+    async def query_impl(_msg, **_kw):
+        # Simulate the agent capturing a validated plan via the tool during
+        # the query (set AFTER the explorer's entry-time capture clear), then
+        # ending with prose that does NOT parse into a sketch.
+        tool_context.solved_plan = grounded_plan
+        tool_context.solved_sketch = captured_sketch
+        return _assistant_response("Solved it. Plan: 1. pick 2. place. Done.")
+
+    explorer._agent_session.query = query_impl
+
+    policy, term_fn = explorer._get_exploration_strategy(0, timeout=5)
+
+    # Recovered (not random fallback): subgoals/options come from the capture.
+    assert callable(policy)
+    assert term_fn(_make_state()) is False
+    assert tool_context.last_sketch_options == [
+        ("Pick", ["block0"]),
+        ("Place", ["block0", "block1"]),
+    ]
+    # The capture was consumed (cleared) so it can't leak into a later solve.
+    assert tool_context.solved_plan is None
+    assert tool_context.solved_sketch is None
+    # Captured params were seeded as initial_params: the option model is
+    # invoked with them (Pick tries them first; Place's are pooled).
+    called = [
+        c.args[1]
+        for c in option_model.get_next_state_and_num_actions.call_args_list
+    ]
+    pick_calls = [o for o in called if o.name == "Pick"]
+    place_calls = [o for o in called if o.name == "Place"]
+    assert pick_calls and place_calls
+    np.testing.assert_allclose(pick_calls[0].params, pick_params)
+    assert any(
+        np.allclose(o.params, place_params) for o in place_calls), \
+        "captured Place params were not seeded into refinement"
+
+
+def test_captured_params_seed_info_gain_search():
+    """With info-seeking ON, the recovered capture's continuous params are
+    seeded as candidates in the info-gain pool (not replayed verbatim)."""
+    _reset_config(agent_explorer_info_seeking=True,
+                  agent_explorer_info_n_feasible_target=2,
+                  agent_bilevel_explorer_max_samples_per_step=4)
+
+    goal_state = _make_state({_block0: [0.5, 0.6, 0.0]})
+    option_model = MagicMock()
+    option_model.get_next_state_and_num_actions.return_value = (goal_state, 3)
+
+    place_params = [0.33, 0.44]
+    grounded_plan, captured_sketch = _make_captured([0.42], place_params)
+
+    explorer, tool_context = _make_explorer(option_model, None)
+    # Wire a trivial ensemble scorer so info-seeking engages on annotated
+    # steps; constant score means the seeded candidate is chosen.
+    tool_context.atom_disagreement_fn = lambda _s, _atoms: 0.0
+
+    async def query_impl(_msg, **_kw):
+        tool_context.solved_plan = grounded_plan
+        tool_context.solved_sketch = captured_sketch
+        return _assistant_response("Done — summary only, no sketch block.")
+
+    explorer._agent_session.query = query_impl
+
+    policy, _ = explorer._get_exploration_strategy(0, timeout=5)
+    assert callable(policy)
+    # Place is the subgoal-annotated step that info-seeking pools; its captured
+    # params must appear among the candidates the pool evaluated.
+    place_calls = [
+        c.args[1]
+        for c in option_model.get_next_state_and_num_actions.call_args_list
+        if c.args[1].name == "Place"
+    ]
+    assert any(np.allclose(o.params, place_params) for o in place_calls), \
+        "captured Place params were not seeded into the info-gain pool"
+
+
 def test_fallback_when_query_fails_and_flag_on():
     """Agent raises → random options fallback when flag enabled."""
     _reset_config(agent_explorer_fallback_to_random=True)