Agent-bilevel: horizon-faithful plan capture + explorer fixes (solves plain domino 5/5 across 5 seeds)

predicators/agent_sdk/bilevel_sketch.py

@@ -1028,6 +1028,15 @@ class ForwardResult:
     first_failure_idx: Optional[int]
     # First step whose positive subgoal atoms diverged, or None.
     first_subgoal_divergence_idx: Optional[int]
+    # Total low-level actions executed across all options.
+    total_actions: int = 0
+    # First step index whose post-state satisfies the goal, or None if the
+    # goal was never reached along the way.
+    goal_step_idx: Optional[int] = None
+    # Cumulative low-level actions through ``goal_step_idx`` (the count the
+    # real, horizon-capped executor would spend to first reach the goal), or
+    # None if the goal was never reached.
+    actions_to_goal: Optional[int] = None
 
     @property
     def executed_all(self) -> bool:
@@ -1039,6 +1048,25 @@ def success(self) -> bool:
         """True iff every option executed AND the goal was reached."""
         return self.executed_all and self.goal_reached
 
+    @property
+    def clean_to_goal(self) -> bool:
+        """True iff the goal is reached with no hard option failure at or
+        before the step that first reaches it.
+
+        Mirrors the real closed-loop executor, which aborts at the first
+        failing option: a 0-action / not-initiable failure *after* the
+        goal already holds is harmless, but one before it dooms the
+        rollout. ``execute_plan_forward`` itself continues past such
+        failures (the option model returns an unchanged post-state), so
+        this guards against capturing a plan whose goal atoms only hold
+        because forward simulation pressed on through a collision the
+        real env would abort on.
+        """
+        if not self.goal_reached or self.goal_step_idx is None:
+            return False
+        return (self.first_failure_idx is None
+                or self.first_failure_idx > self.goal_step_idx)
+
 
 def execute_plan_forward(
     task: Task,
@@ -1048,6 +1076,7 @@ def execute_plan_forward(
     predicates: Set[Predicate],
     sketch: Optional[List[SketchStep]] = None,
     on_step: Optional[Callable[[int, StepOutcome], None]] = None,
+    stop_on_failure: bool = False,
 ) -> ForwardResult:
     """Execute a fully-grounded plan step by step through the option model.
 
@@ -1059,16 +1088,24 @@ def execute_plan_forward(
     ``get_next_state_and_num_actions`` (catching
     ``EnvironmentFailure``), treat 0 actions as a failure, and — when a
     sketch is given — check the step's positive ``subgoal_atoms``
-    against the post-state. Execution stops early only when a step is
-    not initiable or raises (no post-state to continue from); a 0-action
-    step is recorded as a failure but execution continues from the
-    model's returned state. ``on_step(i, outcome)`` is called after each
+    against the post-state. ``on_step(i, outcome)`` is called after each
     step for callers that emit per-step reporting.
+
+    Execution always stops early when a step is not initiable or raises
+    (no post-state to continue from). When ``stop_on_failure`` is True it
+    *also* stops at a 0-action step — mirroring the real closed-loop
+    executor, which aborts at the first failing option rather than
+    pressing on. When False (the default), a 0-action step is recorded as
+    a failure but execution continues from the model's returned state
+    (kept for ``validate_plan_forward``'s per-step diagnostics).
     """
     state = task.init
     steps: List[StepOutcome] = []
     first_failure_idx: Optional[int] = None
     first_div_idx: Optional[int] = None
+    total_actions = 0
+    goal_step_idx: Optional[int] = None
+    actions_to_goal: Optional[int] = None
 
     for i, option in enumerate(plan):
         pre = state
@@ -1121,14 +1158,27 @@ def execute_plan_forward(
             first_failure_idx = i
         if post is None:
             break  # cannot continue without a post-state
+        if failure_reason is not None and stop_on_failure:
+            break  # mirror the real executor: abort at the first failure
         state = post
+        total_actions += num_actions
+        # Record when the goal *first* holds, plus the cumulative actions to
+        # get there — the budget the real horizon-capped executor would
+        # spend. A plan whose goal only holds after more steps than the
+        # episode horizon allows is not real-executable.
+        if goal_step_idx is None and task.goal_holds(state):
+            goal_step_idx = i
+            actions_to_goal = total_actions
 
     return ForwardResult(
         steps=steps,
         final_state=state,
         goal_reached=task.goal_holds(state),
         first_failure_idx=first_failure_idx,
         first_subgoal_divergence_idx=first_div_idx,
+        total_actions=total_actions,
+        goal_step_idx=goal_step_idx,
+        actions_to_goal=actions_to_goal,
     )
 
 

predicators/agent_sdk/docker_agent_runner.py

@@ -80,6 +80,10 @@ async def _run_query(query_input: Dict[str, Any]) -> Dict[str, Any]:
         "type": "enabled",
         "budget_tokens": thinking_budget
     } if thinking_budget and thinking_budget > 0 else None)
+    # Reasoning effort: one of the SDK levels, or unset for ""/"default".
+    effort = str(query_input.get("reasoning_effort", "")).strip().lower()
+    reasoning_effort = (effort if effort in {"low", "medium", "high", "max"}
+                        else None)
     options = ClaudeAgentOptions(
         allowed_tools=allowed_tools,
         mcp_servers={"predicator_tools": mcp_server},
@@ -88,6 +92,7 @@ async def _run_query(query_input: Dict[str, Any]) -> Dict[str, Any]:
         model=query_input["model_name"],
         max_turns=query_input.get("max_turns", 20),
         thinking=thinking,  # type: ignore[arg-type]
+        effort=reasoning_effort,  # type: ignore[arg-type]
     )
 
     client = ClaudeSDKClient(options=options)

predicators/agent_sdk/docker_sandbox.py

@@ -255,6 +255,8 @@ async def query(self,
                 "system_prompt": self._system_prompt,
                 "model_name": self._model_name,
                 "max_turns": CFG.agent_sdk_max_agent_turns_per_iteration,
+                "thinking_budget_tokens": CFG.agent_sdk_thinking_budget_tokens,
+                "reasoning_effort": CFG.agent_sdk_reasoning_effort,
                 "tool_names": self._tool_names,
                 "cfg_snapshot": dict(CFG.__dict__),
                 "log_path": container_log_path,

predicators/agent_sdk/local_sandbox.py

@@ -190,6 +190,15 @@ async def start_session(self) -> None:
             "type": "enabled",
             "budget_tokens": CFG.agent_sdk_thinking_budget_tokens
         } if CFG.agent_sdk_thinking_budget_tokens > 0 else None)
+        # Reasoning effort: pass through to the agent when set to one of the
+        # SDK's accepted levels; "" / "default" leaves it unset.
+        effort = CFG.agent_sdk_reasoning_effort.strip().lower()
+        valid_efforts = {"low", "medium", "high", "max"}
+        if effort and effort not in valid_efforts and effort != "default":
+            raise ValueError(f"agent_sdk_reasoning_effort must be one of "
+                             f"{sorted(valid_efforts)} or ''/'default'; got "
+                             f"{CFG.agent_sdk_reasoning_effort!r}")
+        reasoning_effort = effort if effort in valid_efforts else None
         options = ClaudeAgentOptions(
             allowed_tools=allowed_tools,
             mcp_servers={"predicator_tools": mcp_server},
@@ -198,6 +207,7 @@ async def start_session(self) -> None:
             model=self._model_name,
             max_turns=CFG.agent_sdk_max_agent_turns_per_iteration,
             thinking=thinking,  # type: ignore[arg-type]
+            effort=reasoning_effort,  # type: ignore[arg-type]
             cwd=self._sandbox_dir,
             setting_sources=["project", "local"],
             hooks=(extra_hooks

predicators/agent_sdk/tools.py

@@ -1590,17 +1590,32 @@ def _report_step(i: int, outcome: Any) -> None:
             if img_block and img_block.get("saved_path"):
                 saved_image_paths.append(img_block["saved_path"])
 
+        # Execute exactly like the real closed-loop executor: abort at the
+        # first failing option (0-action collision / not-initiable / env
+        # failure) instead of pressing on. Otherwise forward simulation can
+        # continue past a collision and report a goal that the real rollout —
+        # which ends the episode at that failed option — never reaches.
         result = bilevel_sketch.execute_plan_forward(task,
                                                      grounded_plan,
                                                      ctx.option_model,
                                                      predicates=all_predicates,
                                                      sketch=sketch_steps,
-                                                     on_step=_report_step)
+                                                     on_step=_report_step,
+                                                     stop_on_failure=True)
 
         final_atoms = utils.abstract(result.final_state, ctx.predicates)
         # Use the env's goal-check (its own classifiers); robust to invented
         # predicates that don't reuse env names.
-        goal_achieved = result.goal_reached
+        goal_reached = result.goal_reached
+        # One more real-executor constraint the option model doesn't enforce:
+        # the episode is capped at `CFG.horizon` low-level steps. A plan whose
+        # goal is reached only after more steps than the horizon allows will
+        # time out in real rollout, so don't count it as achieved/captured.
+        horizon = CFG.horizon
+        within_horizon = (result.actions_to_goal is not None
+                          and result.actions_to_goal <= horizon)
+        goal_achieved = (goal_reached and result.clean_to_goal
+                         and within_horizon)
         # Capture a goal-reaching plan on the current task with a sketch that
         # keeps only the subgoals that actually held (so the closed-loop
         # monitor won't flag a spurious divergence on a wrong annotation).
@@ -1649,7 +1664,17 @@ def _report_step(i: int, outcome: Any) -> None:
             goal_str = ", ".join(str(g) for g in sorted(task.goal))
             lines.append(f"Goal: {{{goal_str}}}")
         lines.append(f"Goal achieved: {goal_achieved}")
-        if not goal_achieved and not task.goal_nl:
+        # Goal atoms hold but the plan needs more low-level steps than the
+        # episode horizon allows: say so and that it was NOT captured, so the
+        # agent shortens the plan instead of stopping on a false positive.
+        if goal_reached and not within_horizon:
+            lines.append(
+                f"NOT EXECUTABLE (plan was NOT captured): reaching the goal "
+                f"takes {result.actions_to_goal} low-level steps but the "
+                f"episode horizon is {horizon}. The real executor will run "
+                f"out of steps — shorten the plan (fewer or quicker steps) "
+                f"before resubmitting.")
+        if not goal_reached and not task.goal_nl:
             missing = task.goal - final_atoms
             missing_str = ", ".join(str(a) for a in sorted(missing))
             lines.append(f"Missing goal atoms: {{{missing_str}}}")

predicators/approaches/agent_bilevel_approach.py

@@ -23,6 +23,7 @@
 from predicators import utils
 from predicators.agent_sdk import bilevel_sketch
 from predicators.agent_sdk.bilevel_sketch import SketchStep as _SketchStep
+from predicators.agent_sdk.tools import BUILTIN_TOOLS
 from predicators.approaches import ApproachFailure
 from predicators.approaches.agent_planner_approach import AgentPlannerApproach
 from predicators.execution_monitoring.subgoal_annotations_monitor import \
@@ -215,13 +216,31 @@ def _build_solve_prompt(self,
             all_predicates=self._get_all_predicates(),
             all_options=self._get_all_options(),
             trajectory_summary=self._build_trajectory_summary(),
-            tool_names=self._get_solve_tool_names(),
+            tool_names=self._solve_prompt_tool_names(),
             prior_failures=failures_text,
             initial_image_section=self._initial_image_section(),
             propose_params=CFG.agent_bilevel_use_llm_initial_params,
             require_tool_validation=not CFG.agent_bilevel_refine_fallback,
         )
 
+    def _solve_prompt_tool_names(self) -> Optional[List[str]]:
+        """Tool list advertised in the solve prompt's "Available Tools".
+
+        Mirrors what the explore prompt lists (the explorer renders
+        ``agent_session.tool_names``): the same MCP subset *plus* the
+        sandbox's built-in tools (Bash/Read/Write/...). The built-ins are
+        only actually granted under the local or docker sandbox -- which
+        is exactly when ``LocalSandboxSessionManager.tool_names`` prepends
+        them -- so they are advertised only then. Without a sandbox the
+        list is the bare MCP subset, unchanged.
+        """
+        names = self._get_solve_tool_names()
+        if names is None:
+            return None
+        if CFG.agent_sdk_use_local_sandbox or CFG.agent_sdk_use_docker_sandbox:
+            return list(BUILTIN_TOOLS) + names
+        return names
+
     # ------------------------------------------------------------------ #
     # Solving
     # ------------------------------------------------------------------ #

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
@@ -64,6 +65,30 @@ def _get_exploration_strategy(self, train_task_idx: int,
         # falls back to random before producing one.
         self._tool_context.last_mental_model_solved = None
 
+        # Point the agent's interactive tools (refine_plan_sketch,
+        # evaluate_option_plan, visualize_state) at the EXPLORE task. These
+        # tools default to ctx.current_task when the agent omits task_idx,
+        # and the test-time _solve leaves current_task set to the last TEST
+        # 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.
+        #
+        # 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 = True
+        self._tool_context.solved_plan = None
+        self._tool_context.solved_sketch = None
+
         try:
             prompt = bilevel_sketch.build_solve_prompt(
                 task,
@@ -93,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")
 
@@ -228,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]:

predicators/settings.py

@@ -1034,6 +1034,12 @@ class GlobalSettings:
     # "exceeded the output token maximum" overflow. 0 ⇒ leave unset (the
     # model's default adaptive thinking, which can overflow on hard tasks).
     agent_sdk_thinking_budget_tokens = 16000
+    # Reasoning effort for the agent SDK's Claude agent. One of "low",
+    # "medium", "high", "max" to set it, or "" / "default" to leave it unset
+    # (the model's own default). Higher effort trades latency/cost for more
+    # deliberate planning; orthogonal to agent_sdk_thinking_budget_tokens,
+    # which caps a single response's extended thinking.
+    agent_sdk_reasoning_effort = ""
     agent_sdk_agent_timeout = 300  # seconds per iteration
     agent_sdk_resume_session = True  # resume previous session if available
     agent_sdk_propose_types = True

scripts/configs/predicatorv3/envs/all.yaml

@@ -15,7 +15,7 @@ ENVS:
       excluded_objects_in_state_str: "loc,rot,angle,direction"
       # excluded_predicates is set per-approach: agents.yaml excludes
       # these (test ours); oracle.yaml leaves them in (test oracle).
-      horizon: 400
+      horizon: 500
       domino_initialize_at_finished_state: False
       domino_use_domino_blocks_as_target: True
       domino_use_continuous_place: True