diff --git a/predicators/agent_sdk/belief_probe.py b/predicators/agent_sdk/belief_probe.py index 1a0defd3b3..2f50b52108 100644 --- a/predicators/agent_sdk/belief_probe.py +++ b/predicators/agent_sdk/belief_probe.py @@ -1205,20 +1205,6 @@ def predicates(self, max_trajectories=max_trajectories, max_groundings_per_predicate=max_groundings_per_predicate) - def samplers(self) -> str: - """Reload ``samplers.py`` and install its per-skill samplers. - - Sampler-synthesis sessions only. Loads ``LEARNED_SAMPLERS`` - fresh from the file (snapshotting it into - ``samplers_versions/``), validates the option-name -> callable - map, installs it so ``refine`` draws from the draft samplers, - and reports a per-option sanity check (return shape, in-box - draws) on a representative train-task state. Call it after every - edit of ``samplers.py``. - """ - self._require_available("samplers") - return self._artifact_loader("samplers")() - def _artifact_loader(self, name: str) -> Callable[..., str]: ctx = self._ctx _check_time_budget(ctx) @@ -2903,7 +2889,6 @@ def _on_rollout() -> None: rng=rng, max_draws=max(1, int(max_draws)), top_k=max(1, int(top_k)), - parameterized_samplers=ctx.parameterized_samplers or None, on_rollout=_on_rollout) return ProbeSuggestResult(suggestions, list(notices) + notes) @@ -2974,7 +2959,6 @@ def plan_scorer(plan: List[Any], rng=rng, max_draws=max(1, max_draws), top_k=max(1, int(top_k)), - parameterized_samplers=ctx.parameterized_samplers or None, on_rollout=lambda: _check_time_budget(ctx), plan_scorer=plan_scorer) notices.append( @@ -3095,7 +3079,6 @@ def gated_solved_check(states: List[State], labels: List[Any], check_subgoals=True, check_final_goal=require_goal, run_id="probe", - parameterized_samplers=ctx.parameterized_samplers or None, strip_latent_wait_targets=not ctx.latent_tracking_available, solved_check=solved_check) refined_plan, success = outcome.plan, outcome.success @@ -3219,7 +3202,6 @@ def _select_on_joint_draws( check_subgoals=True, check_final_goal=require_goal, run_id="probe", - parameterized_samplers=ctx.parameterized_samplers or None, strip_latent_wait_targets=not ctx.latent_tracking_available, solved_check=solved_check) extra_samples += outcome.total_samples diff --git a/predicators/agent_sdk/config.py b/predicators/agent_sdk/config.py index 9930b66cdd..65eb8a0da3 100644 --- a/predicators/agent_sdk/config.py +++ b/predicators/agent_sdk/config.py @@ -29,9 +29,7 @@ class SessionConfig: max_turns: int max_buffer_size: int agent_timeout: int - use_docker_sandbox: bool use_local_sandbox: bool - docker_image: str use_scratchpad: bool @classmethod @@ -44,9 +42,7 @@ def from_cfg(cls) -> "SessionConfig": max_turns=CFG.agent_sdk_max_agent_turns_per_iteration, max_buffer_size=CFG.agent_sdk_max_buffer_size, agent_timeout=CFG.agent_sdk_agent_timeout, - use_docker_sandbox=CFG.agent_sdk_use_docker_sandbox, use_local_sandbox=CFG.agent_sdk_use_local_sandbox, - docker_image=CFG.agent_sdk_docker_image, use_scratchpad=CFG.agent_planner_use_scratchpad, ) diff --git a/predicators/agent_sdk/docker_agent_runner.py b/predicators/agent_sdk/docker_agent_runner.py deleted file mode 100644 index 560795a273..0000000000 --- a/predicators/agent_sdk/docker_agent_runner.py +++ /dev/null @@ -1,331 +0,0 @@ -"""Agent runner for Docker sandbox. - -Executed inside the Docker container by DockerSessionManager. Loads a -pickled ``QueryInput`` dict, creates a ``ClaudeSDKClient`` session with -both Claude built-in tools (Bash, Read, Write, Edit, Glob, Grep, Task*) -and custom predicator MCP tools, queries the agent, and pickles results -back to a shared directory. - -The predicators source tree is mounted read-only at ``/opt/predicators`` -(via ``PYTHONPATH``) for imports. Curated reference files are available -at ``/sandbox/reference/``. A writable sandbox is at ``/sandbox``. -PreToolUse hooks restrict the agent's built-in tools to ``/sandbox/``. - -Usage (inside Docker):: - - PYTHONPATH=/opt/predicators python3 \ - /opt/predicators/predicators/agent_sdk/docker_agent_runner.py \ - /data/query_input.pkl /data/query_output.pkl -""" -import asyncio -import logging -import sys -import traceback -from typing import Any, Dict, List, Optional - -import dill as pkl - -# Bootstrap: import predicators.utils before anything else so that Python -# resolves the circular import chain (structs → utils → image_patch_wrapper -# → structs) in the correct order. Without this, importing predicators.structs -# first causes image_patch_wrapper to try "from predicators.structs import Mask" -# while structs is still being initialized, raising an ImportError. -import predicators.utils # noqa: F401, E402 # pylint: disable=unused-import - -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", -) -logger = logging.getLogger(__name__) - -# pylint: disable=wrong-import-position -from predicators.agent_sdk.log_formatter import \ - format_conversation_markdown # noqa: E402 -from predicators.agent_sdk.session_base import build_agent_options, \ - build_sandbox_mcp, stream_agent_response # noqa: E402 - - -async def _run_query(query_input: Dict[str, Any]) -> Dict[str, Any]: - """Create a ClaudeSDKClient, query the agent, and collect responses.""" - from claude_agent_sdk import \ - ClaudeSDKClient # pylint: disable=import-outside-toplevel - - ctx = query_input["tool_context"] - tool_names: Optional[List[str]] = query_input.get("tool_names") - - # MCP server and options come from the same helpers the host-side - # managers use; every value is an explicit query_input entry (no CFG - # reads in-container). An invalid reasoning_effort raises here and - # surfaces as an error response, matching host-side validation. - mcp_server, allowed_tools = build_sandbox_mcp(ctx, tool_names) - options = build_agent_options( - system_prompt=query_input["system_prompt"], - model_name=query_input["model_name"], - allowed_tools=allowed_tools, - mcp_server=mcp_server, - max_turns=query_input.get("max_turns", 20), - # Sent by DockerSessionManager from its SessionConfig; the 20MB - # fallback only covers pickles from older hosts. - max_buffer_size=query_input.get("max_buffer_size", 20 * 1024 * 1024), - reasoning_effort=str(query_input.get("reasoning_effort", "")), - ) - - client = ClaudeSDKClient(options=options) - await client.connect() - - # Incremental log file path (on shared /data or /log volume) - log_path = query_input.get("log_path") - log_meta = {"query": query_input.get("message", "")} - - def _flush_log(collected: List[Dict[str, Any]]) -> None: - """Write current conversation state as markdown to the log file.""" - if not log_path: - return - try: - content = format_conversation_markdown(collected, - title="Docker Query", - meta=log_meta) - with open(log_path, "w", encoding="utf-8") as lf: - lf.write(content) - except Exception: # pylint: disable=broad-except - pass # Don't let logging errors break the agent - - # Docker-specific stderr reporting for real-time host visibility - # (the host streams container stderr into its own log). - def _report_block(dt: float, preview: str) -> None: - print(f"[+{dt:.2f}s] {preview}", file=sys.stderr, flush=True) - - def _report_result(entry: Dict[str, Any]) -> None: - print( - f"Agent iteration complete. " - f"Turns: {entry.get('num_turns', '?')}, " - f"Cost: ${entry.get('total_cost_usd', '?')}", - file=sys.stderr, - flush=True) - - try: - collected = await stream_agent_response( - client, - query_input["message"], - log_label="Docker runner", - report_block=_report_block, - on_result=_report_result, - flush=_flush_log, - ) - finally: - try: - await client.disconnect() - except Exception: # pylint: disable=broad-except - pass - - return { - "responses": collected, - } - - -def _rehash_objects_after_unpickle(ctx: Any) -> None: - """Fix stale Object hash caches after cross-process unpickling. - - ``Object.__hash__`` returns a ``cached_property`` (``_hash``) that - stores ``hash(str(self))``. Python randomises string hashes across - processes (PYTHONHASHSEED), so cached values from the *pickling* - process are stale here. When the option-model simulator later - creates fresh Objects (e.g. ``self._robot`` in ``_get_state``), - their hashes differ from the unpickled Objects, causing KeyError on - ``State.data`` dict lookups. - - Fix: clear every Object's cached ``_hash`` (and ``_str``) so it is - re-computed with the current process's hash seed, then rebuild every - ``State.data`` dict so its internal hash-table is consistent. - """ - from predicators.structs import \ - State # pylint: disable=import-outside-toplevel - - seen: set = set() - - def _clear(obj: Any) -> None: - oid = id(obj) - if oid in seen: - return - seen.add(oid) - obj.__dict__.pop("_hash", None) - obj.__dict__.pop("_str", None) - - def _process_state(state: Any) -> None: - if state is None or not isinstance(state, State): - return - for obj in list(state.data.keys()): - _clear(obj) - # Rebuild dict so Python re-hashes keys with current seed. A - # comprehension (not ``dict(...)``) is load-bearing: ``dict(d)`` - # copies each entry's stored hash without calling ``__hash__``, - # so it would preserve exactly the stale table this repairs. - # pylint: disable-next=unnecessary-comprehension - state.data = {obj: vals for obj, vals in state.data.items()} - - def _process_atoms(atoms: Any) -> None: - for atom in atoms: - for obj in atom.objects: - _clear(obj) - - def _process_task(task: Any) -> None: - # Task has .init (State) and .goal (Set[GroundAtom]) - # EnvironmentTask has .init_obs and .goal_description - if hasattr(task, "init"): - _process_state(task.init) - if hasattr(task, "init_obs"): - _process_state(task.init_obs) - for attr in ("goal", "alt_goal", "goal_description", "alt_goal_desc"): - atoms = getattr(task, attr, None) - # goal_description may be a plain NL string on - # EnvironmentTask; only atom collections carry Objects. - if atoms and not isinstance(atoms, str): - _process_atoms(atoms) - - # Train tasks - for task in getattr(ctx, "train_tasks", []): - _process_task(task) - - # Current task - if ctx.current_task is not None: - _process_task(ctx.current_task) - - # Example state - _process_state(getattr(ctx, "example_state", None)) - - # Trajectories - for traj in (getattr(ctx, "offline_trajectories", []) + - getattr(ctx, "online_trajectories", [])): - for state in traj.states: - _process_state(state) - - -def main() -> None: - """Entry point for Docker agent runner.""" - if len(sys.argv) != 3: - print(f"Usage: {sys.argv[0]} ", - file=sys.stderr) - sys.exit(1) - - input_path = sys.argv[1] - output_path = sys.argv[2] - - logger.info("Docker agent runner starting: input=%s output=%s", input_path, - output_path) - - # Load query input - with open(input_path, "rb") as f: - query_input = pkl.load(f) - - # Restore host CFG settings (arg-specific settings like - # max_num_steps_option_rollout are not set by default import) - if "cfg_snapshot" in query_input: - from predicators.settings import \ - CFG # pylint: disable=import-outside-toplevel - for k, v in query_input["cfg_snapshot"].items(): - setattr(CFG, k, v) - - # Fix stale Object hash caches from cross-process pickling. - ctx = query_input.get("tool_context") - if ctx is not None: - _rehash_objects_after_unpickle(ctx) - - # Recreate option model — the simulator (e.g. PyBullet physics - # server) is process-local and cannot survive pickling. - if ctx is not None and ctx.option_model is not None: - from predicators.option_model import \ - create_option_model # pylint: disable=import-outside-toplevel - from predicators.settings import \ - CFG as _cfg # pylint: disable=import-outside-toplevel - logger.info("Recreating option model (%s) inside Docker...", - _cfg.option_model_name) - ctx.option_model = create_option_model( - _cfg.option_model_name, - skip_residual_dynamics=_cfg.agent_planner_use_base_simulator) - # Sync with all options in context (GT + any previously proposed) - # after the model has its physics server set up. - ctx.option_model._name_to_parameterized_option = { # pylint: disable=protected-access - o.name: o - for o in ctx.options - } - - # Recreate SkillConfig in skill_factory_context — the robot's - # physics_client_id is process-local and stale after pickling. - if (ctx is not None - and ctx.skill_factory_context.get("skill_config") is not None): - from predicators.settings import \ - CFG as _cfg # pylint: disable=import-outside-toplevel - if _cfg.env.startswith("pybullet"): - try: - # pylint: disable=import-outside-toplevel,reimported - from predicators import utils as _utils - from predicators.envs.base_env import BaseEnv - from predicators.envs.pybullet_env import PyBulletEnv - from predicators.ground_truth_models.skill_factories import \ - SkillConfig - - # Find the PyBulletEnv subclass (envs already imported above - # by create_option_model → create_new_env). - env_cls = None - for cls in _utils.get_all_subclasses(BaseEnv): - if (not cls.__abstractmethods__ - and issubclass(cls, PyBulletEnv) - and cls.get_name() == _cfg.env): - env_cls = cls - break - - if env_cls is None: - logger.warning( - "Could not find PyBulletEnv for %s; " - "skill_config NOT recreated", _cfg.env) - else: - _, robot, _ = env_cls.initialize_pybullet(using_gui=False) - ctx.skill_factory_context["skill_config"] = SkillConfig( - robot=robot, - open_fingers_joint=robot.open_fingers, - closed_fingers_joint=robot.closed_fingers, - fingers_state_to_joint=( - env_cls._fingers_state_to_joint), # pylint: disable=protected-access - max_vel_norm=_cfg.pybullet_max_vel_norm, - ik_validate=_cfg.pybullet_ik_validate, - robot_init_tilt=getattr(env_cls, 'robot_init_tilt', - 0.0), - robot_init_wrist=getattr(env_cls, 'robot_init_wrist', - 0.0), - ) - logger.info( - "Recreated SkillConfig inside Docker for %s " - "(physics_client_id=%d)", _cfg.env, - robot.physics_client_id) - except Exception as e: # pylint: disable=broad-except - logger.error("Failed to recreate SkillConfig in Docker: %s", - e, - exc_info=True) - - logger.info("Loaded query input: message length=%d, model=%s", - len(query_input.get("message", "")), - query_input.get("model_name", "?")) - - # Run the query - try: - query_output = asyncio.run(_run_query(query_input)) - except Exception as e: # pylint: disable=broad-except - logger.error("Fatal error in agent runner: %s\n%s", e, - traceback.format_exc()) - query_output = { - "responses": [{ - "type": "error", - "error": str(e) - }], - } - - # Save output - with open(output_path, "wb") as f: - pkl.dump(query_output, f) - - logger.info("Docker agent runner finished: %d responses", - len(query_output.get("responses", []))) - - -if __name__ == "__main__": - main() diff --git a/predicators/agent_sdk/docker_sandbox.py b/predicators/agent_sdk/docker_sandbox.py deleted file mode 100644 index 69ec6dd8c6..0000000000 --- a/predicators/agent_sdk/docker_sandbox.py +++ /dev/null @@ -1,499 +0,0 @@ -"""Docker-sandboxed agent session manager. - -Runs ``ClaudeSDKClient`` inside a Docker container so that the agent's -built-in tools (Bash, Read, Write, Edit, Glob, Grep, Task*) all execute -in an isolated environment. Custom predicator MCP tools are created in-process -inside the container via the same ``create_mcp_tools()`` code used on -the host. - -The host predicators source tree is mounted read-only at -``/opt/predicators`` for Python imports (``PYTHONPATH``). PreToolUse -hooks block the agent's built-in tools (Read, Write, Edit, Glob, Grep) -from accessing anything outside ``/sandbox/``, so the agent cannot -browse environment source code or ground truth models directly. Curated -reference files are copied into ``/sandbox/reference/`` for the agent to -read. The agent can write and run Python scripts in ``/sandbox/``, and -``from predicators.structs import State`` works via the mount. - -Shared data (pickled context and results) passes through ``/data``. - -Behavioral notes relative to the shared base -(:mod:`predicators.agent_sdk.session_base`): - -- ``query()`` is a subprocess orchestrator: each call runs one fresh - container (no persistent client), so ``start_session``, ``close``, - and ``_recover_session`` are no-ops. -- The incremental markdown log is written in-container; the host only - prepends a metadata header afterwards. -- Cost accounting reuses the base delta scheme with the baseline reset - to zero per query, since every container session starts from zero. - -Usage ------ -When the ``agent_sdk_use_docker_sandbox`` flag is ``True``, the -``AgentSessionMixin`` creates a ``DockerSessionManager`` in place of the -normal ``AgentSessionManager``. The interface is identical:: - - manager = DockerSessionManager(...) - responses = await manager.query("Solve this task...") - await manager.close() - -Build the image first:: - - bash docker/build.sh -""" -import datetime -import json -import logging -import os -import shutil -import subprocess -import sys -import tempfile -import uuid -from pathlib import Path -from typing import Any, Dict, List, Optional - -import dill as pkl - -from predicators.agent_sdk.config import SessionConfig -from predicators.agent_sdk.sandbox_prompts import build_sandbox_system_prompt -from predicators.agent_sdk.session_base import SandboxSessionManagerBase -from predicators.agent_sdk.tools import ToolContext, session_log_filename -from predicators.settings import CFG - -logger = logging.getLogger(__name__) - -# Grace period past the per-query agent timeout before the container is -# force-killed (covers container startup + result pickling). -_CONTAINER_TIMEOUT_SLACK_S = 120 - -# Tail sizes for error reporting when a container run fails. -_STDIO_TAIL_CHARS = 2000 -_STDERR_TAIL_LINES = 20 - -# Build Docker-specific prompts from shared templates. -# CLAUDE.md (sandbox mechanics only; see build_claude_md) is written -# into the sandbox when it is populated. -_SANDBOX_SYSTEM_PROMPT = build_sandbox_system_prompt( - env_description="an isolated Docker sandbox", - workspace_description="/sandbox/", - ref_path="/sandbox/reference/", -) - -# --------------------------------------------------------------------------- -# Helper functions -# --------------------------------------------------------------------------- - - -def _get_claude_oauth_token() -> Optional[str]: - """Extract the Claude Code OAuth access token from the macOS Keychain. - - Returns ``None`` on non-macOS platforms or when the token cannot be - found. On macOS, ``claude login`` stores credentials under the - service name ``"Claude Code-credentials"``. - """ - if sys.platform != "darwin": - return None - try: # type: ignore[unreachable] - result = subprocess.run( - [ - "security", "find-generic-password", "-s", - "Claude Code-credentials", "-w" - ], - capture_output=True, - text=True, - timeout=5, - check=False, - ) - if result.returncode != 0: - return None - creds = json.loads(result.stdout.strip()) - return creds.get("claudeAiOauth", {}).get("accessToken") - except (subprocess.SubprocessError, json.JSONDecodeError, KeyError): - return None - - -# _flush_log stays unimplemented on purpose: logs flush inside the -# container, and a host-side call should fail loudly. -# pylint: disable-next=abstract-method -class DockerSessionManager(SandboxSessionManagerBase): - """Runs ClaudeSDKClient inside Docker with built-in + custom MCP tools. - - Matches the ``AgentSessionManager`` interface so that all agent-based - approaches work unchanged. Each ``query()`` call: - - 1. Serializes ``ToolContext`` + message to pickle in a temp directory. - 2. Runs ``docker run ...`` with the predicators source mounted at - ``/opt/predicators:ro`` (for Python imports) and a curated sandbox - at ``/sandbox`` (for agent file operations). - 3. Inside Docker, the runner script creates ``ClaudeSDKClient`` with - both built-in tools AND custom MCP tools, queries the agent, and - pickles back responses + mutated proposals. - 4. Host reads back the pickled results. - - PreToolUse hooks restrict the agent's built-in tools (Read, Write, - Edit, Glob, Grep) to ``/sandbox/`` only. Python imports via - ``PYTHONPATH`` are unaffected. - """ - - _log_label = "Docker" - - def __init__( - self, - system_prompt: str, - log_dir: str, - model_name: str, - tool_context: ToolContext, - tool_names: Optional[List[str]] = None, - image: str = "predicators-sandbox", - extra_reference_files: Optional[Dict[str, str]] = None, - phase: Optional[str] = None, - config: Optional[SessionConfig] = None, - ) -> None: - # Append sandbox instructions to the system prompt - super().__init__(system_prompt=system_prompt + _SANDBOX_SYSTEM_PROMPT, - log_dir=log_dir, - model_name=model_name, - tool_context=tool_context, - tool_names=tool_names, - extra_reference_files=extra_reference_files, - phase=phase, - config=config) - self._image = image - self._last_kind: str = "query" - - # -- Session lifecycle -- - - async def start_session(self) -> None: - """No-op: each query() is a fresh docker run.""" - - async def close(self) -> None: - """No-op: the sandbox directory is kept on disk for inspection.""" - - async def _recover_session(self) -> None: - """No-op: each query is independent.""" - - async def query(self, - message: str, - kind: str = "query") -> List[Dict[str, Any]]: - """Run the agent in Docker and return collected response messages. - - Returns the same ``List[Dict[str, Any]]`` format as - ``AgentSessionManager.query()``. - """ - self._query_count += 1 - self._tool_context.turn_id = self._query_count - self._last_kind = kind - - # Ensure sandbox is set up (lazy init, persists across queries) - self._ensure_sandbox_dir() - - # 1. Create temp directory for data exchange - tmp_dir = tempfile.mkdtemp(prefix="pred-docker-") - input_path = os.path.join(tmp_dir, "query_input.pkl") - output_path = os.path.join(tmp_dir, "query_output.pkl") - - # Compute final log filename upfront so the container can write - # directly to the log directory (incremental updates visible on host). - # Counter-first layout: alphabetical sort matches chronological - # order across mixed ``learn``/``test``/``explore`` phases. - timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - log_filename = session_log_filename( - self._query_count, kind, timestamp, - getattr(self._tool_context, "test_task_idx", None)) - if self._log_dir: - os.makedirs(self._log_dir, exist_ok=True) - incremental_log_path = os.path.join(self._log_dir, log_filename) - else: - incremental_log_path = os.path.join(tmp_dir, "query_log.md") - - try: - # 2. Pickle QueryInput - # Tell the container where to write the incremental log. - # If _log_dir is set, it's mounted at /log inside the container. - container_log_path = (f"/log/{log_filename}" - if self._log_dir else "/data/query_log.md") - query_input = { - "tool_context": self._tool_context, - "message": message, - "system_prompt": self._system_prompt, - "model_name": self._model_name, - "max_turns": self._config.max_turns, - "max_buffer_size": self._config.max_buffer_size, - "reasoning_effort": self._config.reasoning_effort, - "tool_names": self._tool_names, - "cfg_snapshot": dict(CFG.__dict__), - "log_path": container_log_path, - } - with open(input_path, "wb") as f: - pkl.dump(query_input, f) - - logger.info( - "Docker query %d: message length=%d, model=%s", - self._query_count, - len(message), - self._model_name, - ) - - # 3. Build docker run command. Resolve authentication once - # per query (the Keychain OAuth lookup is a subprocess call - # shared by the command and env builders). - api_key = os.environ.get("ANTHROPIC_API_KEY") - oauth_token = None if api_key else _get_claude_oauth_token() - container_name = f"pred-sandbox-{uuid.uuid4().hex[:8]}" - docker_cmd = self._build_docker_command(container_name, tmp_dir, - api_key, oauth_token) - - # 4. Run Docker container - logger.info( - "Starting Docker sandbox: container=%s image=%s", - container_name, - self._image, - ) - env = self._build_env(api_key, oauth_token) - - proc = subprocess.Popen( - docker_cmd, - env=env, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - ) - - # Stream stderr in real-time so tool calls / agent messages - # appear on the host terminal as they happen. - stderr_lines: List[str] = [] - try: - timeout_sec = (self._config.agent_timeout + - _CONTAINER_TIMEOUT_SLACK_S) - import threading # pylint: disable=import-outside-toplevel - - def _stream_stderr() -> None: - assert proc.stderr is not None - for line in proc.stderr: - line = line.rstrip("\n") - stderr_lines.append(line) - logger.info("%s", line) - - stderr_thread = threading.Thread(target=_stream_stderr, - daemon=True) - stderr_thread.start() - - # Wait for stdout (captured for error reporting) - stdout_data = proc.stdout.read() if proc.stdout else "" - proc.wait(timeout=timeout_sec) - stderr_thread.join(timeout=5) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait() - logger.error("Docker container timed out after %ds", - timeout_sec) - stdout_data = "" - - if proc.returncode != 0: - logger.error( - "Docker container exited with code %d.\nstdout: %s\n" - "stderr (last 2000 chars): %s", - proc.returncode, - stdout_data[-_STDIO_TAIL_CHARS:] - if stdout_data else "(empty)", - "\n".join(stderr_lines)[-_STDIO_TAIL_CHARS:] - if stderr_lines else "(empty)", - ) - else: - logger.info("Docker container exited successfully.") - - # 5. Load query output - if os.path.exists(output_path): - with open(output_path, "rb") as f_in: - query_output = pkl.load(f_in) - - responses = query_output.get("responses", []) - # Track costs/turns via the base delta accounting. Each - # docker query is a fresh in-container session whose - # cumulative cost restarts from zero, so reset the delta - # baseline first: every result then charges its full - # cumulative cost. - self._last_cost_usd = 0.0 - for resp in responses: - if resp.get("type") == "result": - self._account_result(resp) - else: - logger.error( - "No output pickle found at %s. Container may have " - "crashed.", output_path) - responses = [{ - "type": - "error", - "error": - (f"Docker container failed (exit code " - f"{proc.returncode}). " - f"stderr: {''.join(stderr_lines[-_STDERR_TAIL_LINES:])}"), - }] - - # 7. Finalize query log - the incremental log was written - # directly to _log_dir as markdown (updated per-message). - # Prepend host metadata header now that the container is done. - if os.path.exists(incremental_log_path) and self._log_dir: - try: - with open(incremental_log_path, encoding="utf-8") as lf: - existing = lf.read() - header_lines = [ - f"- **Query:** {self._query_count}", - f"- **Timestamp:** {timestamp}", - f"- **Session:** {self._session_id}", - f"- **Image:** {self._image}", - "", - "", - ] - with open(incremental_log_path, "w", - encoding="utf-8") as lf: - lf.write("\n".join(header_lines) + existing) - logger.info("Finalized docker query/response at %s", - incremental_log_path) - except Exception: # pylint: disable=broad-except - logger.warning("Failed to enrich log at %s", - incremental_log_path, - exc_info=True) - else: - self._save_query_response_log(message, responses) - - # Track in-memory for conversation replay - self._conversation_log.append({ - "query": message, - "response": responses, - }) - - self._track_fatal_response(responses) - return responses - - finally: - # Cleanup temp data directory (sandbox persists across queries) - shutil.rmtree(tmp_dir, ignore_errors=True) - - def _session_info_extras(self) -> Dict[str, Any]: - """Extra session-info keys: manager type + container image.""" - return { - "session_type": "docker", - "docker_image": self._image, - } - - # -- Internal helpers -- - - def _build_docker_command(self, container_name: str, tmp_dir: str, - api_key: Optional[str], - oauth_token: Optional[str]) -> List[str]: - """Build the ``docker run`` command.""" - cmd = [ - "docker", - "run", - "--rm", - "--name", - container_name, - "--cap-add=NET_ADMIN", - "--cap-add=NET_RAW", - ] - - # Authentication: prefer ANTHROPIC_API_KEY, fall back to OAuth - if api_key: - cmd += ["-e", "ANTHROPIC_API_KEY"] - elif oauth_token: - # The token value itself is added to env in _build_env() - cmd += ["-e", "CLAUDE_CODE_OAUTH_TOKEN"] - else: - # Fall back to bind-mounting ~/.claude - claude_cfg = Path( - os.environ.get("CLAUDE_CONFIG_DIR", - str(Path.home() / ".claude"))) - cmd += ["-v", f"{claude_cfg}:/home/node/.claude"] - - # Mount predicators source for Python imports (hidden from agent - # tools by the PreToolUse hook - only Python's import system can - # read these files). - cmd += ["-v", f"{self._repo_root}:/opt/predicators:ro"] - cmd += ["-e", "PYTHONPATH=/opt/predicators"] - - # Mount curated sandbox directory - cmd += ["-v", f"{self._sandbox_dir}:/sandbox"] - - # Mount data exchange directory - cmd += ["-v", f"{tmp_dir}:/data"] - - # Mount log directory for incremental log updates visible on host - if self._log_dir: - log_dir_abs = os.path.abspath(self._log_dir) - cmd += ["-v", f"{log_dir_abs}:/log"] - - # Working directory - cmd += ["-w", "/sandbox"] - - # Image - cmd.append(self._image) - - # Command: run the agent runner script from the mounted source - cmd += [ - "python3", - "-u", - "/opt/predicators/predicators/agent_sdk/docker_agent_runner.py", - "/data/query_input.pkl", - "/data/query_output.pkl", - ] - - return cmd - - def _build_env(self, api_key: Optional[str], - oauth_token: Optional[str]) -> Dict[str, str]: - """Build environment dict for the docker subprocess.""" - # Pass through host env, stripping CLAUDECODE* vars - env = { - k: v - for k, v in os.environ.items() if not k.startswith("CLAUDECODE") - } - - # Ensure ANTHROPIC_API_KEY is passed through if set - if api_key: - env["ANTHROPIC_API_KEY"] = api_key - elif oauth_token: - env["CLAUDE_CODE_OAUTH_TOKEN"] = oauth_token - - return env - - def _save_query_response_log(self, query: str, - response: List[Dict[str, Any]]) -> None: - """Save query and response to a timestamped markdown file.""" - if not self._log_dir: - return - - timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - kind = self._last_kind - filename = session_log_filename( - self._query_count, kind, timestamp, - getattr(self._tool_context, "test_task_idx", None)) - filepath = os.path.join(self._log_dir, filename) - - lines = [ - f"- **Query:** {self._query_count}", - f"- **Timestamp:** {timestamp}", - f"- **Session:** {self._session_id}", - f"- **Image:** {self._image}", - "", - "# Docker Query", - "", - "## Prompt", - "", - query, - "", - "## Response", - "", - ] - for entry in response: - lines.append( - f"```json\n{json.dumps(entry, indent=2, default=str)}\n```") - lines.append("") - - os.makedirs(self._log_dir, exist_ok=True) - with open(filepath, "w", encoding="utf-8") as f: - f.write("\n".join(lines)) - - logger.info("Saved docker query/response to %s", filepath) diff --git a/predicators/agent_sdk/learn_prompts.py b/predicators/agent_sdk/learn_prompts.py index 6383757024..75f38a195b 100644 --- a/predicators/agent_sdk/learn_prompts.py +++ b/predicators/agent_sdk/learn_prompts.py @@ -303,87 +303,3 @@ def build_program_learn_message( def render_program_zero_shot_message() -> str: """The no-data note for the program arm's first message.""" return render("learn_program_message", "zero_shot") - - -# --------------------------------------------------------------------------- -# Natural-language world model arm (C3) -# --------------------------------------------------------------------------- - - -def build_notes_learn_system_prompt() -> str: - """Compose the natural-language world-model learn system prompt.""" - return _join([ - render("learn_notes_system", "intro"), - render("learn_notes_system", "produce"), - render("learn_notes_system", "tools"), - render("learn_notes_system", "deliverables"), - render("learn_notes_system", "workflow"), - ]) - - -def render_notes_solve_system_section() -> str: - """The solve / explore system-prompt section naming the document.""" - return render("learn_notes_system", "solve_system") - - -def render_world_model_notes_block(notes: str, notes_path: str) -> str: - """The document, quoted into a task message; empty when no notes.""" - if not notes.strip(): - return "" - return render("learn_notes_system", - "notes_block", - notes_path=notes_path, - notes=notes.strip("\n")) - - -def build_notes_learn_message( - *, - n_trajs: int, - n_transitions: int, - n_demos: int, - n_interaction: int, - trajectory_listing: str, - structs_ref: str, - predicate_listing: str, - types_digest: str, - options_digest: str, - notes_file: str, - goal_nls: Sequence[str] = (), - has_prior_notes: bool = False, - objective_block: str = "", - tools_block: str = "", - extra_messages: Sequence[str] = (), -) -> str: - """Compose the natural-language world-model learn first message.""" - goals = [g for g in dict.fromkeys(goal_nls) if g] - goal_block = (render("learn_notes_message", - "goal", - goals="\n".join(f"- {g}" - for g in goals)) if goals else "") - prior_block = (render("learn_notes_message", - "prior_notes", - notes_file=notes_file) if has_prior_notes else "") - body = render( - "learn_notes_message", - "skeleton", - n_trajs=str(n_trajs), - n_transitions=str(n_transitions), - n_demos=str(n_demos), - n_interaction=str(n_interaction), - trajectory_listing=trajectory_listing.strip("\n"), - objective_block=objective_block, - goal_block=goal_block, - prior_notes_block=prior_block, - structs_ref=structs_ref, - predicate_listing=predicate_listing, - types_digest=types_digest.strip("\n"), - options_digest=options_digest.strip("\n"), - tools_block=tools_block, - notes_file=notes_file, - ) - return _join([body, *extra_messages]) - - -def render_notes_zero_shot_message() -> str: - """The no-data note for the natural-language arm's first message.""" - return render("learn_notes_message", "zero_shot") diff --git a/predicators/agent_sdk/log_formatter.py b/predicators/agent_sdk/log_formatter.py index f63fe5cc12..cdf110e004 100644 --- a/predicators/agent_sdk/log_formatter.py +++ b/predicators/agent_sdk/log_formatter.py @@ -2,8 +2,7 @@ Converts the ``List[Dict[str, Any]]`` collected by response parsers into a human-readable markdown document. Used by -``LocalSandboxSessionManager._flush_log`` and -``docker_agent_runner._flush_log``. +``LocalSandboxSessionManager._flush_log``. """ import json from typing import Any, Dict, List, Optional diff --git a/predicators/agent_sdk/policy_execution.py b/predicators/agent_sdk/policy_execution.py index 416c39a19b..4615ffd30a 100644 --- a/predicators/agent_sdk/policy_execution.py +++ b/predicators/agent_sdk/policy_execution.py @@ -13,10 +13,7 @@ build a fresh instance per episode/rollout. * :func:`execute_policy_forward` - the closed-loop sibling of ``plan_execution.execute_plan_forward``, used for belief-model - validation. The real executor - (``AgentModelBasedApproach._policy_to_execution_policy``) mirrors its - semantics step for step, so validation and real execution share one - behavioral contract: + validation, with this behavioral contract: - OPTION failures (not initiable, env failure, 0 actions) do NOT end the episode: the failure text is surfaced to the policy via diff --git a/predicators/agent_sdk/prompts/learn_notes_message.md b/predicators/agent_sdk/prompts/learn_notes_message.md deleted file mode 100644 index 0f375cb618..0000000000 --- a/predicators/agent_sdk/prompts/learn_notes_message.md +++ /dev/null @@ -1,71 +0,0 @@ -# Natural-language world model (learn) first message - -Composed by `AgentNotesWorldModelApproach._build_notes_learn_message` -through `learn_prompts.build_notes_learn_message`. - - -Write the world model document for this environment. There are -__N_TRAJS__ recorded trajectories (__N_TRANSITIONS__ skill-level -transitions) available: __N_DEMOS__ oracle demonstration(s), which -reached the goal by construction, and __N_INTERACTION__ interaction -trajectory/ies collected during online learning, some of which may -have failed to reach the goal. - -__TRAJECTORY_LISTING__ - -Each trajectory carries a `train_task_idx`. `is_goal_state(state, -task_idx)` (equivalently `train_tasks[task_idx].goal_holds(state)`) -checks a single state for the goal atoms. Use it to confirm which -trajectories reached the goal and to treat failed interaction -trajectories as counterexamples: places where the environment -disagreed with what a skill was expected to do. - -__OBJECTIVE_BLOCK__ - -__GOAL_BLOCK__ - -__PRIOR_NOTES_BLOCK__ - -Data-structure source code is at: __STRUCTS_REF__ - -## Available Predicates - -__PREDICATE_LISTING__ - -## Object Types - -__TYPES_DIGEST__ - -## Options - -__OPTIONS_DIGEST__ - -__TOOLS_BLOCK__ - -## This session - -Read the data-structures file first, then explore the trajectory data -with `run_python`. Write your world model to `__NOTES_FILE__` under the -headings given in the system prompt, and finish with the deliverables -listed there. - - -## Task goals (natural language) - -__GOALS__ - - -A `world_model.md` from an earlier cycle exists at `__NOTES_FILE__`. -Read it first; this cycle's data may confirm, refine, or contradict -what it says. Revise it in place. - - -## Zero-shot synthesis - -No trajectory has been recorded and none will be before you finish: -this session is the whole learning phase, and what you write here is -what the planner reasons with on the test tasks. The trajectory counts -above are zero for that reason. Build the document from the task -description, the object types and options, and your own knowledge of -the mechanisms involved, and label every claim as a hypothesis with -the evidence you would want for it. diff --git a/predicators/agent_sdk/prompts/learn_notes_system.md b/predicators/agent_sdk/prompts/learn_notes_system.md deleted file mode 100644 index 3363f521a9..0000000000 --- a/predicators/agent_sdk/prompts/learn_notes_system.md +++ /dev/null @@ -1,106 +0,0 @@ -# Natural-language world model (learn phase) system prompt - -Composed by `AgentNotesWorldModelApproach._get_agent_system_prompt` -through `learn_prompts.build_notes_learn_system_prompt` when the -approach is in its learning phase. The paper's natural-language -world-model arm: the same loop and experiments as the code arms, but the -model is a text document the agent reasons over, never executable. - - -You are building a world model for a robotic manipulation environment -as a natural-language document. No simulator will run from what you -write: at planning time the same document is all the knowledge of the -environment's dynamics the planner has, and it plans by reasoning over -it, so what you write must let a careful reader predict what every -skill does, when it works, and how the environment's own processes -unfold over time. - - -## What you produce - -One file, `world_model.md` (path given in the first message). Keep it -organized under fixed headings so later cycles and the planner can -find things: - -1. `# Mechanisms`: every process the environment runs on its own - (delayed effects, gradual changes, propagation between objects, - hidden state that changes what skills do), each with its trigger - condition, its rate or duration in low-level steps, what it changes - and by how much, and the evidence (trajectory and step) it comes - from. -2. `# Skills`: for every skill, what it changes in the observed state - when it succeeds (with the numbers: offsets, final poses, feature - values as a function of the parameters), the conditions under - which it fails and what the failure looks like, how many low-level - steps it takes, and which of its continuous parameters matter and - over what ranges. -3. `# Thresholds and geometry`: the quantitative gates the environment - enforces (how close is close enough, which side of a fixture, what - counts as supported), each bracketed by recorded attempts on both - sides where the data allows. -4. `# Hidden state`: what the observation does not show, how it can be - inferred from what it does show and from the history of skills - executed, and how it evolves. -5. `# Recipes`: skill sequences, with parameter values, that the data - shows reaching intermediate goals, and why they work. -6. `# Uncertainties and open questions`: what the data does not - settle, phrased as the experiment that would settle it. - -Write for prediction, not description: a reader must be able to take -a state and a skill call and write down the state after it. Prefer -numbers over adjectives, and say where each number comes from. When -you are unsure, say so and give the range. - - -## Tools - -`run_python` is the one tool over the data: `trajectories` -(`List[LowLevelTrajectory]`; each action's `get_option()` is the skill -that produced it, so the skill-level transitions are the spans between -skill changes), `describe_trajectory(i)`, `train_tasks`, -`is_goal_state(state, task_idx)`, and `np`. Use `Read`, `Write` and -`Edit` on `world_model.md`. - - -## Deliverables of a learning session - -- The document, complete under the six headings above, with every - mechanism the recorded episodes exercised reconciled against what - you wrote before (earlier cycles' notes are yours to revise, not to - append to). -- A decision record at the top: the key modeling commitments, the - evidence behind each, and every hypothesis you kept without direct - evidence, labelled as such. -- `./open_questions.md` with what the next exploration should collect - first, and `./strategy.md` with how you would solve the train task - given what you now know. - - -## Workflow - -1. Explore the data with `run_python`: for each skill, which features - change between its start and its end, under what conditions, and by - how much; for each feature that changes while no skill touches it, - what drives it. -2. `Write` or `Edit` `world_model.md`, one heading at a time, with the - numbers and their evidence. -3. Check every claim against a transition it should predict: pick a - recorded skill call, predict its outcome from your notes alone, - compare. Fix the notes where the prediction is wrong. -4. Finish with the deliverables above. - - -## Your world model - -Your knowledge of this environment's dynamics is the natural-language -document `world_model.md`, written during learning; its content is -included in every task message. It is the only model of the -environment you have: there is no simulator to test plans against. -Predict each step of a plan from the document before committing to -it, use the numbers it records for parameters and timings, and treat -its open questions as risks to plan around. - - -## World model notes (__NOTES_PATH__) - -__NOTES__ diff --git a/predicators/agent_sdk/proposal_exec.py b/predicators/agent_sdk/proposal_exec.py index a671599252..c8cc75333b 100644 --- a/predicators/agent_sdk/proposal_exec.py +++ b/predicators/agent_sdk/proposal_exec.py @@ -31,7 +31,7 @@ def _load_sampler_dict( var_name: str, key_error: Callable[[Any], Optional[str]], ) -> Tuple[Dict[str, Any], List[str], Optional[str]]: - """Shared core of the two sampler loaders. + """Core of the ground-sampler loader. Execs ``code`` and validates its ``var_name`` dict. ``key_error`` returns a skip reason for an invalid key (None = valid). Returns @@ -62,27 +62,6 @@ def _load_sampler_dict( return valid, warnings, None -def load_learned_samplers( - code: str, - context: Dict[str, Any], - option_names: Set[str], -) -> Tuple[Dict[str, Any], List[str], Optional[str]]: - """Exec sampler code and validate its ``LEARNED_SAMPLERS`` dict. - - The single loader behind both ``sim.samplers()`` and - ``SamplerLearningMixin._load_samplers_from_module_file``, so the - two cannot drift. Keys must be known option names. - """ - - def key_error(name: Any) -> Optional[str]: - if name not in option_names: - return (f"Skipped '{name}' (not a known option name; known: " - f"{', '.join(sorted(option_names))}).") - return None - - return _load_sampler_dict(code, context, "LEARNED_SAMPLERS", key_error) - - def load_ground_samplers( code: str, context: Dict[str, Any], diff --git a/predicators/agent_sdk/response_parser.py b/predicators/agent_sdk/response_parser.py index c12d52cf84..4dac4be672 100644 --- a/predicators/agent_sdk/response_parser.py +++ b/predicators/agent_sdk/response_parser.py @@ -3,8 +3,8 @@ Converts ``claude_agent_sdk`` message types (``AssistantMessage``, ``UserMessage``, ``ResultMessage``, and the ``SystemMessage`` that marks a context compaction) into plain dicts suitable for logging and -serialization. Used by ``AgentSessionManager``, -``LocalSandboxSessionManager``, and ``docker_agent_runner``. +serialization. Used by ``AgentSessionManager`` and +``LocalSandboxSessionManager``. """ from typing import Any, Dict, List, Optional, Sequence diff --git a/predicators/agent_sdk/sandbox_setup.py b/predicators/agent_sdk/sandbox_setup.py index e41abbfe3a..cebec25855 100644 --- a/predicators/agent_sdk/sandbox_setup.py +++ b/predicators/agent_sdk/sandbox_setup.py @@ -208,7 +208,7 @@ def deny(reason): # this vetoes the interpreter's own import and open events, so a script # cannot reach the hidden predicators modules, their source, or the # harness's run artifacts through computed paths or ``importlib``. Still -# best effort: OS-level isolation (the docker sandbox) is the hard line. +# best effort: it is not OS-level isolation. # --------------------------------------------------------------------------- PYGUARD_DIRNAME = os.path.join(".claude", "pyguard") diff --git a/predicators/agent_sdk/sketch_refinement.py b/predicators/agent_sdk/sketch_refinement.py index ea72664c75..5513445915 100644 --- a/predicators/agent_sdk/sketch_refinement.py +++ b/predicators/agent_sdk/sketch_refinement.py @@ -25,8 +25,8 @@ from predicators.agent_sdk.sketch_types import SketchStep from predicators.option_model import _OptionModelBase from predicators.planning import run_backtracking_refinement -from predicators.structs import GroundAtom, ParameterizedOption, \ - ParameterizedSampler, Predicate, State, Task, _Option +from predicators.structs import GroundAtom, ParameterizedOption, Predicate, \ + State, Task, _Option # Signature of an info-gain scorer: given a candidate post-state and the # atoms whose truth the step is meant to establish, return a scalar where @@ -152,7 +152,6 @@ class _RefineContext: deepest_failure_holder: Optional[List[DeepestFailure]] info_scorer: Optional[InfoScorer] info_n_feasible_target: int - parameterized_samplers: Optional[Dict[str, ParameterizedSampler]] solved_check: Optional[Callable[[List[State], List[Any], bool], Tuple[bool, str]]] # Proposed continuous params are decisions, not seeds: a step that @@ -202,7 +201,6 @@ class _RefinementState: # Options whose synthesized sampler already misbehaved once - so the # per-draw fallback warning fires at most once per option, not on every # one of the (potentially thousands of) draws during backtracking. - sampler_warned: Set[str] = dataclasses.field(default_factory=set) # Step indices whose LLM-proposed initial_params have already been used -- # tried directly on the plain path, or seeded into the info-seeking pool. @@ -248,47 +246,18 @@ class _RefinementState: elapsed: List[float] = dataclasses.field(default_factory=list) -def _draw_params(search: _RefinementState, ctx: _RefineContext, - step: SketchStep, state: State, +def _draw_params(step: SketchStep, state: State, rng_: np.random.Generator) -> np.ndarray: - """Draw continuous params for a step's option. - - Precedence, most specific first: the step's ground sampler (a ``~`` - annotation compiled into a ``GroundSampler``: uniform window or - named code fn), then the option's learned parameterized sampler - (keyed by option name), then uniform ``sample_params`` - the - fallback also on a sampler error or wrong-shaped return (a - misbehaving ground fn falls all the way to uniform, not to the - parameterized sampler, mirroring the parameterized fallback). - """ + """Draw continuous params for a step's option: from the step's ground + sampler (a ``~`` annotation compiled into a ``GroundSampler``: uniform + window or named code fn) when it has one and it draws, uniformly with + ``sample_params`` otherwise.""" if step.ground_sampler is not None: drawn = step.ground_sampler.draw(state, rng_, step.option.params_space, step.objects, step.subgoal_atoms or set()) if drawn is not None: return drawn - return sample_params(step.option, rng_) - sampler = (ctx.parameterized_samplers.get(step.option.name) - if ctx.parameterized_samplers else None) - if sampler is not None: - box = step.option.params_space - expected = box.shape[0] - try: - raw = sampler(state, step.subgoal_atoms or set(), rng_, - list(step.objects)) - params = np.asarray(raw, dtype=np.float32).reshape(-1) - if params.shape == (expected, ): - return np.clip(params, box.low, box.high) - reason = (f"returned shape {params.shape}, " - f"expected ({expected},)") - except Exception as e: # pylint: disable=broad-except - reason = f"raised {type(e).__name__}: {e}" - if step.option.name not in search.sampler_warned: - search.sampler_warned.add(step.option.name) - logging.warning( - "[%s] synthesized sampler for %s %s; falling back to " - "uniform sampling for this option.", ctx.run_id, - step.option.name, reason) return sample_params(step.option, rng_) @@ -347,21 +316,16 @@ def _is_pinned(ctx: _RefineContext, step: SketchStep) -> bool: and step.option.params_space.shape[0] > 0) -def _is_deterministic(ctx: _RefineContext, step: SketchStep) -> bool: - """Whether the step's sampler flags itself as returning constant params.""" - # A sampler may flag itself as returning constant params (ignoring - # state/rng); re-drawing it yields the identical option, so its step - # gets a single attempt -- backtracking then skips straight past it - # instead of wasting the full budget re-descending through it. - if step.ground_sampler is not None: - # A ground-sampler step bypasses the parameterized sampler, - # so a deterministic sampler flag must not collapse it to - # one attempt. An all-zero window pins every draw to the - # center, which IS deterministic - one attempt suffices. - return step.ground_sampler.deterministic - sampler = (ctx.parameterized_samplers.get(step.option.name) - if ctx.parameterized_samplers else None) - return bool(getattr(sampler, "deterministic", False)) +def _is_deterministic(step: SketchStep) -> bool: + """Whether the step's ground sampler returns constant params. + + Re-drawing it yields the identical option, so its step gets a single + attempt: backtracking then skips straight past it instead of wasting + the full budget re-descending through it. An all-zero window pins + every draw to the center, which is deterministic. + """ + return (step.ground_sampler is not None + and step.ground_sampler.deterministic) def _sample_info_seeking(search: _RefinementState, ctx: _RefineContext, @@ -493,8 +457,7 @@ def _consider(grounded: _Option) -> None: if len(scored) > n_pooled_before else "infeasible — not pooled") while len(scored) < ctx.info_n_feasible_target and n_draws < draw_cap: - grounded = ground_step(step, - _draw_params(search, ctx, step, state, rng_)) + grounded = ground_step(step, _draw_params(step, state, rng_)) n_draws += 1 _consider(grounded) pool.spent += n_draws @@ -573,7 +536,7 @@ def _sample_step(search: _RefinementState, ctx: _RefineContext, idx: int, logging.debug("[%s] step %d %s: trying LLM-proposed params %s", ctx.run_id, idx, step.option.name, params.tolist()) return ground_step(step, params) - return ground_step(step, _draw_params(search, ctx, step, state, rng_)) + return ground_step(step, _draw_params(step, state, rng_)) def _validate_step(search: _RefinementState, ctx: _RefineContext, idx: int, @@ -703,7 +666,6 @@ def refine_sketch( deepest_failure_holder: Optional[List[DeepestFailure]] = None, info_scorer: Optional[InfoScorer] = None, info_n_feasible_target: int = 1, - parameterized_samplers: Optional[Dict[str, ParameterizedSampler]] = None, strip_latent_wait_targets: bool = True, solved_check: Optional[Callable[[List[State], List[Any], bool], Tuple[bool, str]]] = None, @@ -798,17 +760,9 @@ def refine_sketch( that ``WaitOption`` terminates on the intended atom change rather than the first incidental one. - ``parameterized_samplers`` maps an option name to a parameterized - (per-skill) sampler ``(state, subgoal_atoms, rng, objects) -> - params`` (the NSRTSampler signature, with the step subgoal in the - atoms slot), used on both plain and info-seeking draws to aim that - option's parameters at the subgoal instead of drawing uniformly. - The return is clipped to the option's box; a missing or misbehaving - sampler falls back to uniform sampling. A step whose sketch line - carries a ``~ [widths]`` region annotation bypasses the sampler - entirely: after the one-shot center try, its draws come from the - step's ``GroundSampler``, the most specific prior winning - ground - sampler, then parameterized sampler, then uniform. + A step whose sketch line carries a ``~ [widths]`` region annotation + draws, after the one-shot center try, from the step's + ``GroundSampler``; other steps draw uniformly. """ if not sketch: return RefineOutcome(plan=[], @@ -833,7 +787,6 @@ def refine_sketch( deepest_failure_holder=deepest_failure_holder, info_scorer=info_scorer, info_n_feasible_target=info_n_feasible_target, - parameterized_samplers=parameterized_samplers, solved_check=solved_check, pin_proposed_params=pin_proposed_params, pinned_step_retries=max(1, pinned_step_retries)) @@ -853,7 +806,7 @@ def refine_sketch( for _step in sketch: if _step.option.params_space.shape[0] == 0: max_tries.append(1) - elif _is_deterministic(ctx, _step): + elif _is_deterministic(_step): max_tries.append(1) elif _is_pinned(ctx, _step): max_tries.append(ctx.pinned_step_retries) @@ -1045,7 +998,6 @@ def suggest_probes( rng: np.random.Generator, max_draws: int = 20, top_k: int = 3, - parameterized_samplers: Optional[Dict[str, ParameterizedSampler]] = None, on_rollout: Optional[Callable[[], None]] = None, plan_scorer: Optional[Callable[[List[_Option], Set[GroundAtom]], Tuple[float, Dict[str, float]]]] = None, @@ -1084,26 +1036,6 @@ def _score(option: _Option, nxt: State, for a in sorted(atoms, key=str) } - ctx = _RefineContext(task=task, - sketch=sketch, - option_model=option_model, - predicates=predicates, - max_samples_per_step=max_draws, - check_subgoals=True, - check_final_goal=False, - log_state=False, - run_id="suggest_probes", - on_step_fail=None, - deepest_failure_holder=None, - info_scorer=info_scorer, - info_n_feasible_target=1, - parameterized_samplers=parameterized_samplers, - solved_check=None, - pin_proposed_params=True, - pinned_step_retries=1) - search = _RefinementState(step_pools=[None] * len(sketch), - step_trajs=[None] * len(sketch), - step_samples_cumulative=[0] * len(sketch)) suggestions: List[StepProbeSuggestion] = [] notes: List[str] = [] state = task.init @@ -1147,8 +1079,7 @@ def _roll(grounded: _Option) -> Optional[State]: best_option: Optional[_Option] = None if has_params and atoms: for _ in range(max_draws): - grounded = ground_step( - step, _draw_params(search, ctx, step, state, rng)) + grounded = ground_step(step, _draw_params(step, state, rng)) n_draws += 1 nxt = _roll(grounded) if nxt is None or not atoms.issubset( @@ -1232,7 +1163,6 @@ def refine_and_validate_report( max_samples_per_step: int, check_subgoals: bool, log_state: bool = False, - parameterized_samplers: Optional[Dict[str, ParameterizedSampler]] = None, run_id: str = "refine", timeout_source: str = "explicit", extra_summary_lines: Optional[List[str]] = None, @@ -1278,7 +1208,6 @@ def refine_and_validate_report( check_subgoals=check_subgoals, log_state=log_state, run_id=run_id, - parameterized_samplers=parameterized_samplers, solved_check=solved_check, strip_latent_wait_targets=strip_latent_wait_targets, ) diff --git a/predicators/agent_sdk/sketch_types.py b/predicators/agent_sdk/sketch_types.py index 113c624572..4cba67acb3 100644 --- a/predicators/agent_sdk/sketch_types.py +++ b/predicators/agent_sdk/sketch_types.py @@ -21,15 +21,11 @@ class GroundSampler: """Per-step (ground) sampler compiled from a sketch annotation. - The ground level of the two-level sampler hierarchy that - ``_draw_params`` consults: ground sampler (this, most specific) > - learned parameterized sampler (``parameterized_samplers``, keyed by - option name) > uniform. A parameterized sampler is authored once - and sees every ground call of its option; a ground sampler is - declared inline for ONE step of ONE sketch and dies with the call - - it lives on the ``SketchStep`` rather than in the option-name-keyed - registry, which could not hold different distributions for two - same-option steps in one sketch. + ``_draw_params`` draws a step's params from its ground sampler when + it has one, uniformly otherwise. A ground sampler is declared inline + for ONE step of ONE sketch and dies with the call: it lives on the + ``SketchStep``, so two same-option steps in one sketch can draw from + different distributions. Two kinds, one per instance: - window (``center`` + ``width`` set): the uniform box a @@ -37,9 +33,9 @@ class GroundSampler: proposed params; - code (``fn`` + ``name`` set): an agent-written function that a ``~ my_sampler`` annotation references by name (loaded fresh per - refine call from the sandbox's ``GROUND_SAMPLERS``); it shares - the parameterized-sampler call signature, so it can shape any - state-conditioned distribution. + refine call from the sandbox's ``GROUND_SAMPLERS``) with the + signature ``(state, subgoal_atoms, rng, objects) -> params``, so + it can shape any state-conditioned distribution. """ center: Optional[np.ndarray] = None width: Optional[np.ndarray] = None @@ -58,7 +54,7 @@ def draw(self, state: State, rng: np.random.Generator, box: Box, Returns ``None`` when a code fn misbehaves (raises or returns a wrong-shaped array); the caller falls back to uniform sampling - for that draw, mirroring the parameterized-sampler fallback. + for that draw. """ if self.fn is not None: try: diff --git a/predicators/agent_sdk/synthesis_backend.py b/predicators/agent_sdk/synthesis_backend.py index 7acc85fbd9..5e84e594bb 100644 --- a/predicators/agent_sdk/synthesis_backend.py +++ b/predicators/agent_sdk/synthesis_backend.py @@ -2,8 +2,8 @@ :class:`SynthesisBackend` declares exactly the approach surface that the synthesis tool factories in :mod:`predicators.agent_sdk.tools` -(``create_synthesis_tools``, ``make_predicate_quality_loader``, -``make_sampler_loader``) and the approach-layer validation +(``create_synthesis_tools``, ``make_predicate_quality_loader``) and +the approach-layer validation glue in :mod:`predicators.approaches.synthesis_validation` dereference. It exists so those modules can be typed against the contract instead of importing the concrete ``AgentSimLearningApproach`` - the import that @@ -22,8 +22,7 @@ from predicators.code_sim_learning.utils import LearnedSimulator from predicators.option_model import _OracleOptionModel from predicators.structs import Action, LowLevelTrajectory, \ - ParameterizedOption, ParameterizedSampler, Predicate, State, Task, \ - Type + ParameterizedOption, Predicate, State, Task, Type class SynthesisBackend(Protocol): @@ -60,8 +59,6 @@ class SynthesisBackend(Protocol): float]]]] # ── State written by the tools ─────────────────────────────── - # Per-skill samplers keyed by option name. - _synthesized_samplers: Dict[str, ParameterizedSampler] # Candidate simulator state published during validation so the # recurrent combined simulator sees the rules under evaluation. _residual_rules: Optional[List] @@ -94,9 +91,6 @@ def _get_all_predicates(self) -> Set[Predicate]: def _get_all_options(self) -> Set[ParameterizedOption]: ... - def _get_all_samplers(self) -> Dict[str, ParameterizedSampler]: - ... - def _group_triples_by_trajectory( self, triples: List[Tuple[State, Action, State]], @@ -137,29 +131,6 @@ def previous_fit_evidence( self, version_tag: str) -> Optional[Dict[str, LaplaceEvidence]]: """The previous version's recorded evidence, as a one-entry dict.""" - def _record_sysid_diagnostics(self, report: Dict[str, Dict[str, Any]], - physical_names: Sequence[str], - num_survivors: int, num_segments: int, - rms: List[float]) -> None: - ... - - def _fit_parameters_recurrent( - self, - rules: List, - specs: List[ParamSpec], - base_pred_triples: List[Tuple[State, Action, State]], - residual_features: Dict[str, List[str]], - ) -> Tuple[FitResult, float]: - ... - - def _fit_parameters_joint_rollout( - self, - rules: List, - rule_specs: List[ParamSpec], - residual_features: Dict[str, List[str]], - ) -> Tuple[FitResult, float]: - ... - def _build_combined_simulator( self, learned_simulator: LearnedSimulator, @@ -195,23 +166,3 @@ class PredicateSynthesisBackend(SynthesisBackend, Protocol): # Initial predicates that survived retraction, used to build the # exec namespace the agent's predicate code runs in. _kept_initial_predicates: Set[Predicate] - - -class SamplerSynthesisBackend(Protocol): - """The narrow surface ``make_sampler_loader`` needs. - - ``SamplerLearningMixin`` satisfies this directly (its declared host- - class contract covers every member), so the mixin can pass ``self`` - without seeing the full backend. - """ - - _fitted_params: Dict[str, float] - _train_tasks: List[Task] - _types: Set[Type] - _synthesized_samplers: Dict[str, ParameterizedSampler] - - def _get_all_predicates(self) -> Set[Predicate]: - ... - - def _get_all_options(self) -> Set[ParameterizedOption]: - ... diff --git a/predicators/agent_sdk/tools/__init__.py b/predicators/agent_sdk/tools/__init__.py index 6bdd8e555e..55258bbd18 100644 --- a/predicators/agent_sdk/tools/__init__.py +++ b/predicators/agent_sdk/tools/__init__.py @@ -17,8 +17,7 @@ - ``python_exec``: shared python-exec core (run_python / run_python). - ``synthesis`` / ``params_view``: the synthesis-session tool factory. -- ``predicate_synthesis`` / ``sampler_synthesis``: the loaders behind - ``sim.predicates()`` / ``sim.samplers()``. +- ``predicate_synthesis``: the loader behind ``sim.predicates()``. This facade re-exports the package's public surface (plus a few underscore names kept for pre-split imports); new code should import @@ -36,7 +35,6 @@ list_session_tool_names from predicators.agent_sdk.tools.results import _make_coercing_tool, \ _make_spilling_text_result, session_log_filename -from predicators.agent_sdk.tools.sampler_synthesis import make_sampler_loader from predicators.agent_sdk.tools.sandbox_guard import \ SANDBOX_HIDDEN_MODULES_PATTERN, SANDBOX_INTROSPECTION, \ SANDBOX_SYSTEM_ROOTS, _screen_text_for_sandbox_escape @@ -73,7 +71,6 @@ "list_session_tool_names", "load_ground_sampler_fns", "make_predicate_quality_loader", - "make_sampler_loader", "make_solved_check", "make_write_snapshot_hook", "render_pybullet_image", diff --git a/predicators/agent_sdk/tools/context.py b/predicators/agent_sdk/tools/context.py index e91941e226..26acb4ab90 100644 --- a/predicators/agent_sdk/tools/context.py +++ b/predicators/agent_sdk/tools/context.py @@ -9,7 +9,7 @@ from predicators.option_model import _OptionModelBase from predicators.settings import CFG from predicators.structs import CausalProcess, LowLevelTrajectory, \ - ParameterizedOption, ParameterizedSampler, Predicate, State, Task, Type + ParameterizedOption, Predicate, State, Task, Type @dataclass(frozen=True) @@ -69,10 +69,9 @@ class ToolContext: # need the real engine). A comparison arm lists what its surface # withholds; the probe raises on a listed call instead of serving it. probe_disabled: FrozenSet[str] = frozenset() - # Synthesis-session loaders behind ``sim.predicates()`` and - # ``sim.samplers()``: each reloads the agent-authored file fresh - # (predicates.py / samplers.py), installs the result into the - # approach so refinement sees the draft, and returns the report + # Synthesis-session loader behind ``sim.predicates()``: it reloads + # the agent-authored predicates.py fresh, installs the result into + # the approach so refinement sees the draft, and returns the report # text. Empty in sessions that do not offer the artifact. probe_artifact_loaders: Dict[str, Callable[..., @@ -111,18 +110,11 @@ class ToolContext: probe_score_provider: Optional[Callable[..., str]] = None # Active-experiment info-gain scorer, synced from the learning # approach when info-seeking exploration is on: - # ``(state, atoms) -> disagreement``. The agent_model_based explorer - # passes it into refinement so continuous-parameter search prefers - # candidates that straddle the learned model's decision boundaries. - # None ⇒ plain feasibility search (default). + # ``(state, atoms) -> disagreement``. The probe passes it into + # refinement so continuous-parameter search prefers candidates that + # straddle the learned model's decision boundaries. None ⇒ plain + # feasibility search (default). atom_disagreement_fn: Optional[Callable[[State, Any], float]] = None - # Synthesized per-skill samplers (option name -> sampler), synced from - # the learning approach when agent_sim_learn_parameterized_samplers is on. - # The agent_model_based explorer and synthesis tools pass these into - # refinement so continuous-parameter search aims at each step's subgoal - # instead of drawing uniformly. Empty ⇒ uniform sampling (default). - parameterized_samplers: Dict[str, ParameterizedSampler] = field( - default_factory=dict) current_task: Optional[Task] = None # The last real observation of the level in progress (continual # play: every env tool result and the session query refresh it), so @@ -153,8 +145,6 @@ class ToolContext: joint_draw_scope: Optional[Callable[[Dict[str, float]], Any]] = None episode_prefix_provider: Optional[Callable[[], Tuple[List[State], List[Any]]]] = None - skill_factory_context: Dict[str, Any] = field(default_factory=dict) - proposals_disabled: bool = False # set True during test-time solving log_dir: Optional[str] = None env: Optional[Any] = None # simulator env reference (for rendering) image_save_dir: Optional[str] = None # sandbox path for rendered images @@ -184,44 +174,15 @@ class ToolContext: # frozen for the session's lifetime. Subclasses set this before # opening a fresh session and clear it on close. extra_session_hooks: Dict[str, list] = field(default_factory=dict) - # Populated by AgentModelBasedExplorer so learning approaches can diff - # mental-model subgoals against real trajectories. - # TODO(sim-learning): consume these in learn_from_interaction_results. - last_sketch_subgoals: Optional[Any] = None - # Agent-session phase the tools are serving: "explore", "solve" - # (test-time) or "synthesis" (learn). Set by the session mixin when - # it builds the session; None before any session exists. + # Agent-session phase the tools are serving: "solve" or + # "synthesis" (the model-writing rounds). Set by the session mixin + # when it builds the session; None before any session exists. phase: Optional[str] = None # True when the approach will track the simulator's latent block at # execution (code_sim_learning.latent_tracker), so latent-reading # atoms are evaluable on real observations and refinement must keep # them as Wait targets; False (bare observations) strips them. latent_tracking_available: bool = False - last_sketch_options: Optional[Any] = None - # Set by AgentModelBasedExplorer per request: did the mental model reach - # the task goal during refinement? Read by get_interaction_requests to - # stamp InteractionRequest.mental_model_solved (None ⇒ no verdict). - last_mental_model_solved: Optional[bool] = None - # Sketch-line descriptions of the exploration plans already generated - # this online-learning cycle (a cycle's requests are all generated - # before any executes). Cleared by get_interaction_requests per cycle, - # appended by AgentModelBasedExplorer per request, and shown in the next - # explore prompt so the agent proposes a complementary plan instead of - # repeating the identical one for every request. - cycle_scheduled_plans: List[str] = field(default_factory=list) - # Digest of the latest rollout system-ID fit's weak spots - # (unexplainable segments, unidentified/insensitive params, - # cross-cycle conflicts), synced from the sim-learning approach. - # The agent_model_based explorer appends it to its experiment guidance - # so the next exploration targets the gaps. None ⇒ no fit ran yet - # (or it had no weak spots). - sysid_diagnostics: Optional[str] = None - # The natural-language world-model arm's document (world_model.md - # content) and its agent-visible path: the solve prompt and the - # model-free explorer quote it into every task message. Empty - # everywhere else. - world_model_notes: str = "" - world_model_notes_path: str = "" # Set by submit_plan / submit_policy when a plan is verified # to reach the goal on the CURRENT solve task: the simulator-verified plan # (grounded options with found params) and the parallel subgoal sketch. @@ -236,9 +197,8 @@ class ToolContext: # below). Cleared together with solved_plan. solved_plan_reached_goal: Optional[bool] = None # Gate for the above: only approaches that consume captured plans - # (AgentModelBasedApproach) set this True. Keeps the open-loop - # planner, which also uses submit_plan, from recording - # spurious captures. + # set this True, so other users of submit_plan record no spurious + # captures. capture_goal_reaching_plans: bool = False # Set (with capture_goal_reaching_plans) only for the final-submission # nudge after an attempt exhausted its turn budget: submit_plan @@ -330,8 +290,8 @@ class ToolContext: # submit_plan keeps its probing role but its CAPTURE gate is # disabled, and submit_policy requires it. Set by _solve_attempt. policy_capture_mode: bool = False - # Restart-loop attempt bookkeeping, set by AgentModelBasedApproach._solve - # around each attempt. ``attempt_start``/``attempt_deadline`` are + # Attempt bookkeeping, set around each attempt (a continual play + # round is one). ``attempt_start``/``attempt_deadline`` are # time.monotonic() values; the deadline is enforced cooperatively by # the probe (every sim call) and run_python, and surfaced in tool # results as a budget footer. None ⇒ no attempt in flight / no wall @@ -418,9 +378,7 @@ def begin_attempt(self, index: int, wall_clock: float) -> None: Resets everything scoped to a single attempt (rollout count, best refused submission) and arms the wall-clock deadline - (``wall_clock <= 0`` ⇒ no deadline). The matching teardown stays - in ``AgentModelBasedApproach._solve``'s finally block, - interleaved with its journal write. + (``wall_clock <= 0`` ⇒ no deadline). """ self.attempt_index = index self.attempt_rollout_count = 0 diff --git a/predicators/agent_sdk/tools/sampler_synthesis.py b/predicators/agent_sdk/tools/sampler_synthesis.py deleted file mode 100644 index 0cbbc231bc..0000000000 --- a/predicators/agent_sdk/tools/sampler_synthesis.py +++ /dev/null @@ -1,172 +0,0 @@ -"""The ``sim.samplers()`` loader for sampler-synthesis sessions.""" -from typing import Any, Callable, Dict, List, Optional, Tuple - -import numpy as np - -from predicators.agent_sdk.proposal_exec import build_exec_context, \ - load_learned_samplers -from predicators.agent_sdk.synthesis_backend import SamplerSynthesisBackend -from predicators.agent_sdk.tools.params_view import _ParamsView -from predicators.agent_sdk.tools.sandbox_guard import _scrub_host_paths -from predicators.agent_sdk.tools.snapshots import _ArtifactSnapshotter -from predicators.settings import CFG -from predicators.structs import Object - - -def make_sampler_loader( - samplers_file: str, - samplers_versions_dir: str, - approach: SamplerSynthesisBackend, - cycle_index_provider: Optional[Callable[[], int]] = None, -) -> Callable[[], str]: - """Build the ``sim.samplers()`` loader for one synthesis session. - - On each call the loader loads - ``samplers.py`` fresh (snapshotting into ``samplers_versions_dir``), - validates the ``LEARNED_SAMPLERS`` dict (option name -> callable), - installs it into ``approach._synthesized_samplers`` so refinement - uses it, and reports a per-option shape/in-box sanity check. - - Args: - samplers_file: Host path to the agent-edited ``samplers.py``. - samplers_versions_dir: Directory for per-call snapshots. - approach: The ``AgentSimLearningApproach`` instance. - cycle_index_provider: Returns the current 0-based cycle - (negative = the offline pass, rendered as ``offline``). - """ - # pylint: disable=import-outside-toplevel - import traceback # pylint: disable=redefined-outer-name,reimported - - from predicators.code_sim_learning.fit_space import ParamSpec - - # pylint: enable=import-outside-toplevel - _snapshotter = _ArtifactSnapshotter( - live_file=samplers_file, - versions_dir=samplers_versions_dir, - artifact_name="samplers", - cycle_index_provider=cycle_index_provider, - missing_file_hint=("Use Write to create it with " - "LEARNED_SAMPLERS = {\"OptionName\": fn, ...}."), - ) - params_view = _ParamsView(approach._fitted_params) # pylint: disable=protected-access - - def _snapshot_and_load_samplers( - path: str, - ) -> Tuple[Dict[str, Any], Optional[str], Optional[str], List[str]]: - """Snapshot ``path`` then exec it into a fresh namespace. - - Returns ``(samplers, version_tag, error_msg, warnings)``. - Entries keyed by an unknown option name, or whose value is not - callable, are skipped and described in ``warnings``. On success, - mutates ``approach._synthesized_samplers`` to the validated - dict. - """ - raw, version_tag, err = _snapshotter.snapshot(path) - if err is not None: - return {}, None, err, [] - assert raw is not None and version_tag is not None - - ctx = build_exec_context( - types=approach._types, # pylint: disable=protected-access - predicates=approach._get_all_predicates(), # pylint: disable=protected-access - options=approach._get_all_options(), # pylint: disable=protected-access - extra_context={ - "params": params_view, - "ParamSpec": ParamSpec, - }) - option_names = {o.name for o in approach._get_all_options()} # pylint: disable=protected-access - valid, warnings, err = load_learned_samplers(raw.decode("utf-8"), ctx, - option_names) - if err is not None: - return {}, version_tag, (f"[{version_tag}] Error executing " - f"{path}:\n{err}"), [] - - # Mutate approach state so sim.refine / test-time - # refinement draw from the agent's draft samplers. - approach._synthesized_samplers = valid # pylint: disable=protected-access - return valid, version_tag, None, warnings - - def _sanity_check(name: str, fn: Any) -> str: - """Draw a few params from a representative state; report shape/box.""" - # pylint: disable=protected-access - options_by_name = {o.name: o for o in approach._get_all_options()} - opt = options_by_name[name] - train_tasks = approach._train_tasks - if not train_tasks: - return f" {name}: no train task to sanity-check against." - state = train_tasks[0].init - # Pick the first object of each option-arg type present in the state. - objs: List[Object] = [] - for t in opt.types: - match = next((o for o in state if o.type.name == t.name), None) - if match is None: - return (f" {name}: no object of type '{t.name}' in the " - "train-task state to sanity-check against.") - objs.append(match) - box = opt.params_space - expected = box.shape[0] - rng = np.random.default_rng(CFG.seed) - in_box = 0 - n_draws = 3 - for _ in range(n_draws): - try: - raw = fn(state, set(), rng, objs) - arr = np.asarray(raw, dtype=np.float32).reshape(-1) - except Exception: # pylint: disable=broad-except - last = traceback.format_exc().strip().splitlines()[-1] - return (f" {name}: ERROR — sampler raised: {last} " - "(note: this check, and refinement at steps with " - "no subgoal annotation, call the sampler with " - "subgoal_atoms=set(); it must not crash on an " - "empty set — fall back to a default or uniform " - "draw).") - if arr.shape != (expected, ): - return (f" {name}: ERROR — returned shape {arr.shape}, " - f"expected ({expected},).") - if bool(np.all(arr >= box.low - 1e-6)) and \ - bool(np.all(arr <= box.high + 1e-6)): - in_box += 1 - return (f" {name}: OK — {n_draws} draws, {in_box}/{n_draws} " - f"within the params box.") - - def sampler_report() -> str: - """Reload samplers.py, install LEARNED_SAMPLERS, and report the per- - option sanity check.""" - try: - samplers, version_tag, err, warnings = ( - _snapshot_and_load_samplers(samplers_file)) - except Exception: # pylint: disable=broad-except - return (f"Error loading samplers.py:\n" - f"{_scrub_host_paths(traceback.format_exc())}") - - if err is not None: - return err - - prefix = f"[{version_tag}]" - lines = [ - f"{prefix} Sampler report — {len(samplers)} per-skill " - f"sampler(s) installed.", - ] - if warnings: - lines.append("") - lines.append("Warnings (entries skipped during load):") - for w in warnings: - lines.append(f" - {w}") - - if not samplers: - lines.append("") - lines.append("LEARNED_SAMPLERS is empty — add " - "{\"OptionName\": fn} entries to samplers.py.") - return "\n".join(lines) - - lines.append("") - lines.append("Sanity check (representative train-task state):") - for name in sorted(samplers): - lines.append(_sanity_check(name, samplers[name])) - lines.append("") - lines.append("Now call sim.refine with a sketch that " - "uses these options to measure the samples-to-refine " - "improvement.") - return "\n".join(lines) - - return sampler_report diff --git a/predicators/agent_sdk/tools/synthesis.py b/predicators/agent_sdk/tools/synthesis.py index ee53e936a2..a4393f50ca 100644 --- a/predicators/agent_sdk/tools/synthesis.py +++ b/predicators/agent_sdk/tools/synthesis.py @@ -546,9 +546,6 @@ def _evaluate_rollout_fit(rules: list, sse=float(outcome.pre_sse), pinned=True, coverage=(0, len(rollouts))) - if hasattr(approach, "_record_sysid_diagnostics"): - approach._record_sysid_diagnostics( # pylint: disable=protected-access - {}, physical_names, 0, len(rollouts), outcome.traj_rms) rms_str = ", ".join(f"{r:.4g}" for r in outcome.traj_rms) trim_threshold = (CFG.code_sim_learning_rollout_trim_rms_factor * DEFAULT_NOISE_SIGMA) @@ -601,10 +598,6 @@ def _evaluate_rollout_fit(rules: list, num_points=CFG.agent_plan_validation_physics_margin_points) if outcome.belief is None else []), belief=outcome.belief) - if hasattr(approach, "_record_sysid_diagnostics"): - approach._record_sysid_diagnostics( # pylint: disable=protected-access - ident_report, physical_names, outcome.num_survivors, - len(rollouts), outcome.traj_rms) kept_at_init = sorted(n for n in physical_names if applied[n] != fitted[n]) # Like-for-like SSE headline: the % reduction is measured on the diff --git a/predicators/approaches/__init__.py b/predicators/approaches/__init__.py index d1e5300eb2..502e46736b 100644 --- a/predicators/approaches/__init__.py +++ b/predicators/approaches/__init__.py @@ -1,6 +1,5 @@ """Handle creation of approaches.""" -import logging from typing import List, Set from typing import Type as TypingType @@ -16,20 +15,8 @@ # Find the subclasses. utils.import_submodules(__path__, __name__) -# Deprecated CLI aliases from before the agent approaches' model-free / -# model-based rename (2026-08-30); old commands and configs still pass -# them. -_DEPRECATED_APPROACH_ALIASES = { - "agent_planner": "agent_model_free", - "agent_bilevel": "agent_model_based", -} - def _get_approach_cls_from_name(name: str) -> TypingType[BaseApproach]: - if name in _DEPRECATED_APPROACH_ALIASES: - logging.warning("Approach name %r is deprecated; use %r.", name, - _DEPRECATED_APPROACH_ALIASES[name]) - name = _DEPRECATED_APPROACH_ALIASES[name] for cls in utils.get_all_subclasses(BaseApproach): if not cls.__abstractmethods__ and cls.get_name() == name: return cls diff --git a/predicators/approaches/agent_continual_approach.py b/predicators/approaches/agent_continual_approach.py index d4f628425f..a07a43f84e 100644 --- a/predicators/approaches/agent_continual_approach.py +++ b/predicators/approaches/agent_continual_approach.py @@ -293,8 +293,7 @@ def _round_extra_tools(self, session: ProtocolSession) -> List[Any]: self._install_extra_synthesis_surfaces(exec_ns, base_pred_triples, inferred_hint, extra_paths) candidate_provider = self._make_candidate_probe_model_provider( - paths.simulator_file, trajectories, base_pred_triples, - inferred_hint) + paths.simulator_file, trajectories) def probe_model() -> _OptionModelBase: # Before a candidate exists the probe runs the real skill @@ -437,7 +436,7 @@ def _deploy_session_model(self, session: ProtocolSession, inferred_hint: Dict[str, List[str]], paths: Any, extra_paths: Dict[str, str]) -> None: loaded = self._load_synthesis_artifacts(trajectories, inferred_hint, - paths, extra_paths, {}) + paths, extra_paths) if loaded is None: # No loadable simulator.py this session: the prior model, if # any, still stands. diff --git a/predicators/approaches/agent_model_based_approach.py b/predicators/approaches/agent_model_based_approach.py deleted file mode 100644 index 4b6c11676a..0000000000 --- a/predicators/approaches/agent_model_based_approach.py +++ /dev/null @@ -1,1298 +0,0 @@ -"""Agent model-based approach: the agent delivers a simulator-validated plan. - -The agent plans a sequence of parameterized skills with object bindings, -subgoal atoms after each step, and continuous parameters, and must -DELIVER it as an ``submit_plan`` capture on the current task - -nothing it did not validate in the simulator (the model) is ever -executed. A backtracking parameter search remains available to the agent -as a probe method (``sim.refine``) and to mid-episode -suffix replans, but there is no approach-side refinement of unvalidated -sketches. - -Registered under the CLI approach name ``agent_model_based`` -(``agent_bilevel`` is kept as a deprecated alias). - -Example command:: - - python predicators/main.py --env pybullet_domino \ - --approach agent_model_based --seed 0 \ - --num_train_tasks 1 --num_test_tasks 1 \ - --num_online_learning_cycles 1 --explorer agent_model_free -""" -import dataclasses -import hashlib -import logging -import os -import time -from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple - -import numpy as np - -from predicators import utils -from predicators.agent_sdk import bilevel_sketch -from predicators.agent_sdk.session_base import AgentSessionFatalError, \ - query_fatal_error -from predicators.agent_sdk.sketch_types import SketchStep as _SketchStep -from predicators.agent_sdk.tools import BUILTIN_TOOLS, load_ground_sampler_fns -from predicators.approaches import ApproachFailure -from predicators.approaches.agent_model_free_approach import \ - AgentModelFreeApproach -from predicators.execution_monitoring.subgoal_annotations_monitor import \ - SubgoalExecutionStatus -from predicators.settings import CFG -from predicators.structs import Action, GroundAtom, Object, \ - ParameterizedOption, Predicate, State, Task, _Option - -# Fraction of agent_solve_attempt_wall_clock below which the attempt's -# budget counts as spent, not merely close to it: agents that watch the -# [budget] footer wrap up shortly BEFORE the deadline, and a query whose -# remaining tools would only refuse has nothing left to give. Used to -# label why an attempt ended (see _attempt_end_reason). -_SPENT_WALL_FRACTION = 0.2 - -# Cap on the natural-language goal text in a task's journal entry: long -# enough for any real goal_nl, short enough that a runaway goal string -# cannot crowd the journal's entry budget. -_JOURNAL_GOAL_MAX_CHARS = 400 - -# Final-submission nudge (see _nudge_final_submission). It runs once per -# task, on the LAST attempt, so it accepts a plan that falls short of a -# validated solve: the restarts are gone and a partial plan beats -# forfeiting the task. -_FINAL_SUBMIT_NUDGE = ( - "You are out of exploration budget for this attempt. Do NOT explore " - "further. In as few tool calls as possible, submit your single best " - "plan NOW via submit_plan on the current task (omit " - "task_idx), using the best parameters you have already validated. " - "It is captured as your answer even if it does not fully reach the " - "goal or does not score as a solve; then finish.") - -# Policy-mode variant: a best-effort POLICY is genuinely better than a -# best-effort plan - it is closed-loop, so whatever recovery logic it -# carries still applies at execution. -_FINAL_SUBMIT_POLICY_NUDGE = ( - "You are out of exploration budget for this attempt. Do NOT explore " - "further. In as few tool calls as possible, submit your current best " - "./policy.py NOW via submit_policy on the current task. It is " - "captured as your answer even if it does not fully reach the goal or " - "does not score as a solve; then finish.") - - -@dataclasses.dataclass -class _CaptureInfo: - """Metadata of the most recently consumed captured plan. - - Recorded by :meth:`AgentModelBasedApproach._consume_validated_plan` so - the restart loop can distinguish a validated solve (return - immediately) from a best-effort capture (bank it, rank across - attempts by evaluator reward) and journal the plan. - """ - validated: bool - reward: Optional[float] - plan_lines: List[str] - # One-line capture-time validation record (rollout tally, first - # failing step, physics-margin tally), journaled so a later - # fresh-context attempt sees HOW reliable the capture was (e.g. - # "8/10 rollouts ok; first failure: ... step 25 (Place) ...") - # instead of only that it exists. - validation_summary: Optional[str] = None - - -class AgentModelBasedApproach(AgentModelFreeApproach): - """Model-based planning: the agent plans a skeleton with subgoals and - parameters and submits it as a simulator-validated capture. - - Extends AgentModelFreeApproach - reuses agent session, tools, - trajectory management, exploration, save/load. Overrides solving - with the capture-only query loop plus the restart/journal machinery. - """ - - def __init__(self, *args: Any, **kwargs: Any) -> None: - super().__init__(*args, **kwargs) - if CFG.agent_bilevel_max_execution_replans > 0 and \ - CFG.execution_monitor != "subgoal_annotations": - raise ValueError( - "agent_bilevel_max_execution_replans > 0 requires " - "--execution_monitor subgoal_annotations (got " - f"{CFG.execution_monitor!r}): divergence detection lives " - "in the execution monitor, so without it test execution " - "is silently open-loop.") - if CFG.agent_solve_policy_mode: - if CFG.agent_bilevel_max_execution_replans > 0: - raise ValueError( - "agent_solve_policy_mode is mutually exclusive with " - "agent_bilevel_max_execution_replans > 0: the policy " - "OWNS closed-loop recovery (option failures are " - "surfaced to it), so the sketch-divergence replan " - "machinery must be off.") - if not CFG.agent_planner_use_simulator: - raise ValueError( - "agent_solve_policy_mode requires " - "agent_planner_use_simulator: the policy is validated " - "in the belief model before execution.") - # Live status of the currently executing annotated plan, exported - # to the subgoal_annotations execution monitor. None whenever no - # monitored plan is active (exploration, replanning disabled). - self._exec_status: Optional[SubgoalExecutionStatus] = None - # The grounded option plan behind _exec_status, kept so a - # divergence with no refinable suffix can resume the remaining - # not-yet-executed options open-loop (the dispensed policy holds - # them only in its closure). Set/cleared alongside _exec_status. - self._exec_plan: Optional[List[_Option]] = None - # Per-episode replan budget, refreshed by reset_for_new_episode. - self._exec_replans_left = 0 - # Whether the most recent sketch query ended because the agent hit - # agent_sdk_max_agent_turns_per_iteration. Set by - # _query_agent_for_plan_sketch; every query ending is terminal for - # its attempt, so this only labels WHY the attempt ended (see - # _attempt_end_reason) in the logs and the journal. - self._last_sketch_query_hit_turn_cap = False - # Why the last capture-less attempt ended (_attempt_end_reason), - # sampled inside the attempt while its budget signals are still - # armed - _solve clears them before writing the journal entry. - self._last_attempt_end_reason = "" - # Metadata of the last capture consumed by - # _consume_validated_plan; read by the restart loop in _solve. - self._last_capture_info: Optional[_CaptureInfo] = None - # Tasks whose goal + init-state journal entry is already written - # (one context entry per task, at the top of its section). - self._journal_task_context_recorded: Set[Any] = set() - # Snapshot of that set at begin_test_phase: test-task keys are - # rolled back with the journal itself, so a later evaluation - # (whose entries were removed) re-writes its context entries. - self._pre_test_journal_context_keys: Optional[Set[Any]] = None - - @classmethod - def get_name(cls) -> str: - return "agent_model_based" - - # ------------------------------------------------------------------ # - # Execution monitoring (closed-loop test execution) - # ------------------------------------------------------------------ # - - def reset_for_new_episode(self) -> None: - super().reset_for_new_episode() - self._exec_status = None - self._exec_plan = None - self._exec_replans_left = CFG.agent_bilevel_max_execution_replans - # Optionally give each test solve a fresh agent conversation. reset() - # fires once per test task (not on mid-episode replans, which go - # through step()); the next query lazily rebuilds the session with the - # same sandbox + artifacts but empty chat context. Test-phase only, so - # exploration episodes keep their shared session. - if CFG.agent_fresh_session_per_test_task and self._in_test_phase: - self._close_agent_session() - - def get_execution_monitoring_info(self) -> List[Any]: - if self._exec_status is None: - return [] - return [self._exec_status] - - def begin_test_phase(self) -> None: - super().begin_test_phase() - self._pre_test_journal_context_keys = set( - self._journal_task_context_recorded) - - def end_test_phase(self) -> None: - super().end_test_phase() - # The journal rollback removed this evaluation's entries, so its - # task-context dedup keys must go too - the same test tasks are - # re-solved next evaluation and need fresh goal + init entries. - if self._pre_test_journal_context_keys is not None: - self._journal_task_context_recorded = \ - self._pre_test_journal_context_keys - self._pre_test_journal_context_keys = None - - # ------------------------------------------------------------------ # - # Agent session hooks - # ------------------------------------------------------------------ # - - def _get_synthesis_tool_names(self) -> Optional[List[str]]: - """No synthesis phase in this approach - declare an empty set.""" - return [] - - # ------------------------------------------------------------------ # - # System prompt - # ------------------------------------------------------------------ # - - def _get_agent_system_prompt(self) -> str: - # Sessions are per-phase (see _ensure_agent_session: a phase - # change closes and rebuilds the session), so a solve session - # only ever receives solve queries and an explore session only - # explore queries: each phase's system prompt states just its - # own deliverable contract and rules. The query carries the task - # data and the run state (sketch_prompts.build_solve_prompt). - return bilevel_sketch.build_solve_system_prompt( - explore=self._explore_phase, - policy_mode=CFG.agent_solve_policy_mode, - propose_params=CFG.agent_bilevel_use_llm_initial_params, - ground_samplers=CFG.agent_bilevel_ground_samplers, - physics_margin=CFG.agent_plan_validation_physics_margin, - rule_param_margin=CFG.agent_plan_validation_rule_param_margin, - necessity=CFG.agent_plan_validation_necessity, - use_journal=CFG.agent_solve_use_journal, - execute_certified_plan=CFG.agent_explorer_execute_certified_plan, - early_stop_note=bilevel_sketch.build_early_stop_note(), - policy_max_options=CFG.agent_policy_max_options, - policy_max_repeated_failures=CFG. - agent_policy_max_repeated_failures, - policy_max_repeated_noops=CFG.agent_policy_max_repeated_noops, - ) - - # ------------------------------------------------------------------ # - # Solve prompt (no continuous params, subgoal format) - # ------------------------------------------------------------------ # - - def _build_solve_prompt(self, task: Task) -> str: - """Build prompt asking for a plan sketch without continuous params.""" - journal_text = "" - strategy_text = "" - attempts_text = "" - if CFG.agent_solve_use_journal: - # pylint: disable-next=import-outside-toplevel - from predicators.agent_sdk import journal as journal_mod - journal_text = journal_mod.read_journal( - self._tool_context.sandbox_dir) - attempts_text = journal_mod.read_journal( - self._tool_context.sandbox_dir, - filename=journal_mod.ATTEMPTS_FILENAME) - strategy_text = journal_mod.read_strategy( - self._tool_context.sandbox_dir) - return bilevel_sketch.build_solve_prompt( - task, - all_predicates=self._get_all_predicates(), - all_options=self._get_all_options(), - trajectory_summary=self._build_trajectory_summary(), - tool_names=self._solve_prompt_tool_names(), - initial_image_section=self._initial_image_section(), - propose_params=CFG.agent_bilevel_use_llm_initial_params, - require_tool_validation=True, - journal=journal_text, - strategy=strategy_text, - attempts=attempts_text, - ) - - 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 - # ------------------------------------------------------------------ # - - def _solve(self, task: Task, timeout: int) -> Callable[[State], Action]: - replan_policy = self._maybe_replan_from_divergence(task, timeout) - if replan_policy is not None: - return replan_policy - ctx = self._tool_context - self._record_task_context_in_journal(task) - max_attempts = max(1, CFG.agent_solve_max_attempts) - wall_clock = CFG.agent_solve_attempt_wall_clock - # Best best-effort capture across attempts, ranked by evaluator - # reward. A validated (evaluator-solved) capture returns - # immediately; only when no attempt produces one does the best - # banked policy execute for its honest reward. - best_policy: Optional[Callable[[State], Action]] = None - best_reward = -float("inf") - last_failure: Optional[ApproachFailure] = None - for attempt in range(1, max_attempts + 1): - if CFG.agent_solve_fresh_context: - # Fresh conversation per attempt (and per test task): a - # failed attempt's context carries its confidently wrong - # world model (run_20260717_230436 seed1's "hard collision - # boundary" that its identical sibling placed through), so - # a restart is the cheapest de-anchoring mechanism. Curated - # knowledge travels through the solve journal instead. - self._close_agent_session() - ctx.begin_attempt(attempt, wall_clock) - self._last_capture_info = None - policy: Optional[Callable[[State], Action]] = None - unexpected: Optional[Exception] = None - try: - policy = self._solve_attempt(task) - except ApproachFailure as e: - last_failure = e - except AgentSessionFatalError: - # The session backend is unusable (auth/billing/config); - # neither a banked capture nor further restarts can help. - # Re-raise so the run terminates (the finally still runs - # for bookkeeping). - raise - except Exception as e: # pylint: disable=broad-except - # ApproachTimeout is a SIBLING of ApproachFailure (both - # subclass ExceptionWithInfo), and env/SDK errors can - # also escape - none of them may skip the cleanup below, - # and a banked capture from an earlier attempt should - # still execute rather than be forfeited (handled after - # the finally). - unexpected = e - finally: - # Attempt bookkeeping must not outlive the attempt on ANY - # exit path (including KeyboardInterrupt): stale fields - # would append bogus [budget] footers and mislabel journal - # entries in later sessions sharing this ToolContext. - ctx.attempt_deadline = None - # Policy mode is scoped to solve attempts: left armed, it - # would silently disable submit_plan's capture - # gate for the EXPLORER's queries, which deliver sketches - # even in policy-mode configs. - ctx.policy_capture_mode = False - info = self._take_capture_info() - self._record_attempt_in_journal(attempt, max_attempts, policy, - info) - ctx.attempt_start = None - ctx.attempt_index = 0 - if unexpected is not None: - if best_policy is not None: - logging.warning( - "[%s] Solve attempt %d/%d raised %s; executing the " - "banked best-effort capture instead of forfeiting.", - self._run_id, attempt, max_attempts, unexpected) - return best_policy - raise unexpected - if policy is not None and (info is None or info.validated): - # Defensive: every capture path records metadata, so a - # missing record is treated as a validated solve rather - # than banked at unknown reward. - return policy - if policy is not None: - assert info is not None - reward = (info.reward - if info.reward is not None else -float("inf")) - if best_policy is None or reward > best_reward: - best_policy = policy - best_reward = reward - if attempt < max_attempts: - logging.info( - "[%s] Solve attempt %d/%d ended without a validated " - "solve%s; restarting with %s context.", self._run_id, - attempt, max_attempts, " (best-effort capture banked)" - if policy is not None else "", - "fresh" if CFG.agent_solve_fresh_context else "the same") - if best_policy is not None: - logging.info( - "[%s] No validated solve in %d attempt(s); executing the " - "best best-effort capture (evaluator reward %s).", - self._run_id, max_attempts, - f"{best_reward:.2f}" if best_reward > -float("inf") else "n/a") - return best_policy - if last_failure is not None: - raise last_failure - raise ApproachFailure( - f"Bilevel solve produced no captured plan in {max_attempts} " - "attempt(s).") - - def _attempt_wall_spent(self) -> bool: - """Whether the attempt's wall clock is spent (or nearly so). - - True once less than :data:`_SPENT_WALL_FRACTION` of the wall - clock remains, not merely once the deadline passes: agents that - watch the [budget] footer wrap up shortly BEFORE the deadline - (run_20260718_125643 ended an attempt with 2 minutes left), and - an attempt whose remaining tools would only refuse is spent in - every sense that matters. Tool-side refusals keep using the - exact deadline, so the closing minutes still allow submissions. - """ - deadline = self._tool_context.attempt_deadline - if deadline is None: - return False - floor = _SPENT_WALL_FRACTION * CFG.agent_solve_attempt_wall_clock - return time.monotonic() > deadline - floor - - def _attempt_end_reason(self) -> str: - """Why the attempt's single query ended, for logs and journal.""" - if self._last_sketch_query_hit_turn_cap: - return "turn cap" - if self._attempt_wall_spent(): - return "wall clock spent" - return "no submission" - - def _take_capture_info(self) -> Optional[_CaptureInfo]: - """Pop the metadata _consume_validated_plan recorded (or None). - - An accessor rather than a bare attribute read: the attribute is - set as a side effect of _solve_attempt, which mypy cannot see, - so reading it directly right after ``= None`` is flagged - unreachable. - """ - info = self._last_capture_info - self._last_capture_info = None - return info - - def _append_journal_auto_entry(self, header: str, - body_lines: List[str]) -> bool: - """Best-effort append of a harness-written solve-journal entry. - - Callers guard with :meth:`_journal_active`. Returns True on a - successful write; a failed write is logged, never raised - the - journal must not be able to fail a solve. - """ - sandbox_dir = self._tool_context.sandbox_dir - assert self._journal_active() and sandbox_dir - # pylint: disable-next=import-outside-toplevel - from predicators.agent_sdk import journal as journal_mod - try: - journal_mod.append_entry( - sandbox_dir, - header, - "\n".join(body_lines), - max_chars=journal_mod.MAX_AUTO_ENTRY_CHARS, - filename=journal_mod.ATTEMPTS_FILENAME) - except OSError as e: - logging.warning("Journal entry %r failed: %s", header, e) - return False - return True - - def _journal_task_label(self) -> str: - """The journal's label for the task currently being solved. - - Carries the HARNESS cycle number so the attempt log anchors a - canonical numbering: without it agents invented their own cycle - counts (a journal's "cycle 4" was the harness's cycle 1), and - cross-referencing run logs against the journal needed a mental - offset. - """ - idx = self._tool_context.test_task_idx - task_part = f"task {idx}" if idx is not None else "task ?" - return f"cycle {self._tool_context.iteration_id} {task_part}" - - def _record_task_context_in_journal(self, task: Task) -> None: - """Append the task's goal + init-state entry, once per task. - - Written at the START of the task's first attempt so it tops the - task's journal section - above even the agent's own in-attempt - notes - keeping every later entry interpretable: a recorded - plan's geometry is only meaningful relative to its layout. The - init dict uses the exact representation the solve prompt shows - (including its excluded-objects filtering). - """ - if not self._journal_active(): - return - key = (self._tool_context.test_task_idx, id(task)) - if key in self._journal_task_context_recorded: - return - if task.goal_nl: - goal_txt = " ".join(task.goal_nl.split()) - if len(goal_txt) > _JOURNAL_GOAL_MAX_CHARS: - goal_txt = goal_txt[:_JOURNAL_GOAL_MAX_CHARS].rstrip() + "..." - else: - goal_txt = ", ".join(str(a) for a in sorted(task.goal, key=str)) - body = [f"- goal: {goal_txt}", "- initial state features:"] - body.extend(f" {line}" - for line in task.init.dict_str(indent=2).splitlines()) - header = f"{self._journal_task_label()} goal + initial state (auto)" - if self._append_journal_auto_entry(header, body): - self._journal_task_context_recorded.add(key) - - def _record_attempt_in_journal(self, attempt: int, max_attempts: int, - policy: Optional[Any], - info: Optional[_CaptureInfo]) -> None: - """Auto-append this attempt's factual record to the attempt log. - - The harness-written record (outcome, budget spent, captured or - best refused plan) guarantees the essentials of every attempt - are on record even when the agent writes nothing; the agent's - own lessons live in journal.md, which it edits directly. - """ - if not self._journal_active(): - return - ctx = self._tool_context - body = [f"- outcome: {self._attempt_outcome_text(policy, info)}"] - if ctx.attempt_start is not None: - elapsed_min = (time.monotonic() - ctx.attempt_start) / 60.0 - body.append(f"- budget spent: {elapsed_min:.1f} min, " - f"{ctx.attempt_rollout_count} sim rollouts") - if info is not None and info.validation_summary: - body.append(f"- {info.validation_summary}") - if info is not None and info.plan_lines: - body.append("- captured plan:") - body.extend(f" {line}" for line in info.plan_lines) - elif ctx.best_uncaptured_plan_lines: - # Nothing captured, but the attempt's best refused submission - # (evaluator non-solve or flaky) is worth carrying: a later - # attempt - or the final best-effort nudge - can resubmit it - # instead of the work vanishing with the attempt's context. - reward_txt = (f"evaluator reward {ctx.best_uncaptured_reward:.2f}" - if ctx.best_uncaptured_reward is not None else - "no evaluator verdict") - body.append(f"- best refused submission ({reward_txt}, " - "not captured):") - body.extend(f" {line}" for line in ctx.best_uncaptured_plan_lines) - header = (f"{self._journal_task_label()} attempt " - f"{attempt}/{max_attempts} (auto)") - self._append_journal_auto_entry(header, body) - - def _attempt_outcome_text(self, policy: Optional[Any], - info: Optional[_CaptureInfo]) -> str: - """One-line outcome for an attempt's journal record.""" - if policy is None: - # The reason is the one fact a fresh-context restart cannot - # rediscover: it tells the next attempt whether the last one - # ran out of budget or talked itself out of submitting. - if self._last_attempt_end_reason: - return f"no capture ({self._last_attempt_end_reason})" - return "no capture" - if info is None or info.validated: - return "SOLVED (validated capture)" - if info.reward is not None: - return f"best-effort capture (evaluator reward {info.reward:.2f})" - return "best-effort capture (no evaluator verdict)" - - def _solve_attempt(self, task: Task) -> Callable[[State], Action]: - """One full solve attempt: a single agent query on one session. - - The attempt's budgets are the wall clock - (``agent_solve_attempt_wall_clock``) and the query's turn cap; - the only deliverable is an ``submit_plan`` capture - (consumed via :meth:`_consume_validated_plan`). - - However that query ends - a spent budget, an unparseable sketch, - or a session that simply never submitted - the attempt is over. - A second query on the same conversation would re-explore from a - context that already contains whatever went wrong - (run_20260808_113951 queries 004-009: three full-price queries - restating the same "no plan can reach the goal" argument), so the - fresh-context restart is the only retry. See :meth:`_end_attempt` - for the one exception, on the final attempt. - """ - self._sync_tool_context() - self._tool_context.current_task = task - # Let submit_plan record a goal-reaching - # plan on this task into solved_plan/solved_sketch (consumed below). - self._tool_context.capture_goal_reaching_plans = True - # Policy mode: the deliverable is policy.py via submit_policy; - # submit_plan stays available for probing but cannot - # capture. - self._tool_context.policy_capture_mode = CFG.agent_solve_policy_mode - # LLM-free bypass: a prewritten policy.py as the captured - # artifact (smoke tests / debugging the execution path). - if CFG.agent_solve_policy_mode and CFG.agent_policy_file: - with open(CFG.agent_policy_file, "r", encoding="utf-8") as f: - self._tool_context.solved_policy_source = f.read() - self._tool_context.solved_plan_reached_goal = True - policy = self._consume_validated_plan() - assert policy is not None - return policy - # Render the initial state so the agent can see the scene layout. - self._render_initial_state_image(task) - # Whether later fresh-context restarts exist after this attempt; - # decides whether this attempt pays for the final-submission nudge - # (see _end_attempt). attempt_index == 0 (no restart loop in - # flight) behaves like a final attempt. - restarts_remain = (0 < self._tool_context.attempt_index < max( - 1, CFG.agent_solve_max_attempts)) - # Clear any prior capture so we only act on this query's result. - self._tool_context.clear_plan_capture() - self._last_sketch_query_hit_turn_cap = False - self._last_attempt_end_reason = "" - try: - self._query_agent_for_plan_sketch(task) - except AgentSessionFatalError: - raise - except Exception as e: # pylint: disable=broad-except - # The agent may have validated a working plan via - # submit_plan even if its final text didn't parse. - policy = self._consume_validated_plan() - if policy is not None: - return policy - logging.warning("[%s] Solve query failed: %s", self._run_id, e) - else: - # Fast path: the agent already refined + forward-validated - # a plan on this task via submit_plan - return it - # directly instead of re-refining the (possibly different) - # final-text sketch. - policy = self._consume_validated_plan() - if policy is not None: - return policy - # The agent must itself reach a confirmed - # submit_plan capture (consumed above) so we - # never execute a plan it didn't verify. - logging.info("[%s] Query ended without a validated plan.", - self._run_id) - # Sample the end reason before _end_attempt: the nudge suspends - # the attempt deadline, and _solve clears it outright before the - # journal entry is written. - self._last_attempt_end_reason = self._attempt_end_reason() - policy = self._end_attempt(restarts_remain) - if policy is not None: - return policy - raise ApproachFailure("Bilevel solve failed: the attempt's agent " - "query produced no captured plan " - f"({self._last_attempt_end_reason}).") - - def _end_attempt( - self, - restarts_remain: bool) -> Optional[Callable[[State], Action]]: - """End an attempt that produced no capture. - - With later fresh-context restarts remaining, end with NO nudge - (return None): the restart is the retry, and the journal - auto-entry already records the attempt's best refused - submission. On the final attempt the best-effort submission - nudge is the ultimate fallback - return whatever policy it - captures (None when even that yields nothing, giving up on the - task). - """ - if restarts_remain: - return None - return self._nudge_final_submission() - - # ------------------------------------------------------------------ # - # Plan sketch extraction - # ------------------------------------------------------------------ # - - def _query_agent_for_plan_sketch(self, task: Task) -> List[_SketchStep]: - """Query agent for a plan sketch and parse it.""" - sketch_file = CFG.agent_bilevel_plan_sketch_file - if sketch_file: - # An absolute path is used as-is; a bare filename is resolved - # against the configured plan-sketch directory under scripts/. - if os.path.isabs(sketch_file): - filepath = sketch_file - else: - filepath = ( - f"{utils.get_path_to_predicators_root()}/scripts/" - f"{CFG.agent_bilevel_plan_sketch_dir}/{sketch_file}") - with open(filepath, "r", encoding="utf-8") as f: - plan_text = f.read().strip() - logging.info("Loaded plan sketch from file: %s", sketch_file) - else: - prompt = self._build_solve_prompt(task) - responses = self._query_agent_sync(prompt, kind="test") - dead = query_fatal_error(responses) - if dead is not None: - # An outage is not a failed attempt: recording 0/1 here - # would write a bogus eval datapoint. Stop the run; the - # relaunch re-runs this cycle's test. - raise AgentSessionFatalError( - "test query died without the agent doing any work " - f"({dead}); not recording this attempt as a failure.") - # Record cap-exhaustion before parsing: a capped session usually - # has no final text, so the "empty plan text" failure below is - # still attributable to the turn cap by _attempt_end_reason. - self._last_sketch_query_hit_turn_cap = \ - self._responses_hit_turn_cap(responses) - plan_text = self._extract_option_plan_text(responses) - - if not plan_text: - raise ApproachFailure("Agent returned empty plan text.") - - # Tolerant parse of the agent's final text; named `~ my_sampler` - # references resolve against the sandbox's ground_samplers.py (a - # broken file just drops the annotations here - this is the - # best-effort fallback path, not the strict tool path). - gs_fns, gs_err = load_ground_sampler_fns(self._tool_context) - if gs_err is not None: - logging.warning("[%s] %s", self._run_id, gs_err) - sketch = bilevel_sketch.parse_sketch_from_text( - plan_text, - task, - predicates=self._get_all_predicates(), - options=self._get_all_options(), - types=self._types, - parse_continuous_params=CFG.agent_bilevel_use_llm_initial_params, - parse_ground_samplers=CFG.agent_bilevel_ground_samplers, - ground_sampler_fns=gs_fns or None, - ) - - if not sketch: - option_names = sorted(o.name for o in self._get_all_options()) - raise ApproachFailure(f"Parsed empty plan sketch from agent.\n" - f" Plan text:\n{plan_text}\n" - f" Available option names: {option_names}") - - logging.info( - "[%s] Agent produced sketch with %d steps, %d with " - "subgoals.", self._run_id, len(sketch), - sum(1 for s in sketch if s.subgoal_atoms)) - return sketch - - @staticmethod - def _responses_hit_turn_cap(responses: List[Dict[str, Any]]) -> bool: - """Whether a query's response stream ended on the SDK turn cap. - - The SDK reports the cap as result subtype ``error_max_turns``; - the num_turns comparison is a fallback for backends whose result - entries lack the subtype field. - """ - max_turns = CFG.agent_sdk_max_agent_turns_per_iteration - for entry in responses: - if entry.get("type") != "result": - continue - if entry.get("subtype") == "error_max_turns": - return True - num_turns = entry.get("num_turns") - if num_turns is not None and num_turns >= max_turns: - return True - return False - - # ------------------------------------------------------------------ # - # Backtracking refinement (used by mid-episode suffix replans) - # ------------------------------------------------------------------ # - - def _refine_sketch( - self, - task: Task, - sketch: List[_SketchStep], - timeout: float, - attempt: int = 0, - on_step_fail: Optional[Callable[[int, List[Optional[_Option]], str], - None]] = None, - ) -> Tuple[List[_Option], bool]: - """Backtracking search over continuous parameters for a plan sketch. - - Returns ``(plan, success)``. On success, ``plan`` is a list of - grounded options that achieves the task goal. On failure, - ``plan`` is the longest partial refinement found. - - This is the approach-flavored entry to - ``bilevel_sketch.refine_sketch`` (which stays settings-free): - it gathers the approach-owned inputs (option model, predicates, - samplers, run id), reads the search knobs from ``CFG``, and - first passes the task through :meth:`_attach_initial_latent` so - partially-observable approaches can seed ``task.init.latent`` - with the initial latent block. Used by mid-episode suffix - replans. - - ``attempt`` perturbs the RNG so retries explore different - samples - without it, refinement is deterministic in - ``CFG.seed`` and a forward-validation failure would loop on - the identical plan. ``on_step_fail`` is forwarded to the search - (called with the step index, the partial plan, and the failure - reason whenever a step fails to refine). - """ - task = self._attach_initial_latent(task) - assert self._option_model is not None, \ - "agent_bilevel requires a simulator " \ - "(agent_planner_use_simulator=True)." - outcome = bilevel_sketch.refine_sketch( - task, - sketch, - self._option_model, - predicates=self._get_all_predicates(), - timeout=timeout, - rng=np.random.default_rng(CFG.seed + attempt), - max_samples_per_step=CFG.agent_bilevel_max_samples_per_step, - check_subgoals=CFG.agent_bilevel_check_subgoals, - log_state=CFG.agent_bilevel_log_state, - run_id=self._run_id, - parameterized_samplers=self._get_all_samplers(), - on_step_fail=on_step_fail, - strip_latent_wait_targets=( - not self._tool_context.latent_tracking_available), - ) - return outcome.plan, outcome.success - - def _attach_initial_latent(self, task: Task) -> Task: - """Hook for partial-observability approaches to seed the latent. - - Subclasses that thread a ``latent`` state block through the - simulator (``AgentSimLearningApproach`` when the loaded rules - are recurrent) override this to attach an initial latent to - ``task.init.latent`` before refinement begins. The default - returns ``task`` unchanged - fully-observable approaches need do - nothing. - """ - return task - - def _sample_params(self, option: ParameterizedOption, _state: State, - rng: np.random.Generator) -> np.ndarray: - """Sample continuous parameters for an option.""" - return bilevel_sketch.sample_params(option, rng) - - def _parse_subgoal_annotations( - self, - text: str, - predicates: Set[Predicate], - objects: Sequence[Object], - ) -> List[Optional[Tuple[Set[GroundAtom], Set[GroundAtom]]]]: - """Shim over ``bilevel_sketch.parse_subgoal_annotations``.""" - option_names = {o.name for o in self._get_all_options()} - return bilevel_sketch.parse_subgoal_annotations( - text, predicates, objects, option_names) - - # ------------------------------------------------------------------ # - # Helpers - # ------------------------------------------------------------------ # - - def _maybe_replan_from_divergence( - self, task: Task, - timeout: int) -> Optional[Callable[[State], Action]]: - """Handle a mid-episode re-solve triggered by the subgoal_annotations - execution monitor. - - CogMan calls solve() identically at episode start and on a - monitor-triggered replan; ``_exec_status`` distinguishes them - (non-None only while a monitored plan executes; - reset_for_new_episode clears it at episode start). On a replan - ``task.init`` is the real state where the just-finished step's - annotation failed. Divergence is usually a continuous-execution - problem (a sampled parameter whose real outcome differed from - the option-model rollout), not a wrong skeleton, so we first try - to resume a suffix of the executed sketch (cheap, no agent - query; see :meth:`_replan_suffix`). When no suffix refines - or - the episode's replan budget is spent - the remaining - not-yet-executed options resume OPEN-LOOP instead of failing the - episode: an annotation is the agent's prediction, not proof the - goal is out of reach, and aborting a plan whose remaining - settle/cure steps might still deliver turns a maybe-fail into a - certain fail. The divergence stays in the log and the goal check - decides the episode. Set - ``CFG.agent_bilevel_replan_agent_fallback`` to instead fall - through to a fresh agent sketch query when no suffix refines. - """ - status = self._exec_status - if status is None or status.steps_initiated == 0: - return None - self._exec_status = None - exec_plan = self._exec_plan or [] - self._exec_plan = None - failed_idx = status.steps_initiated - 1 - steps = list(status.sketch) - failed_name = steps[failed_idx].option.name - if self._exec_replans_left > 0: - self._exec_replans_left -= 1 - logging.info( - "Subgoal divergence after step %d (%s). Replanning from the " - "current state (%d execution replans left).", failed_idx, - failed_name, self._exec_replans_left) - policy = self._replan_suffix(task.init, task, steps, failed_idx, - timeout) - if policy is not None: - return policy - if CFG.agent_bilevel_replan_agent_fallback: - # No suffix of the executed skeleton refines from here; - # fall through to pay for a fresh agent sketch. - logging.info("Suffix replan failed; querying the agent for " - "a fresh sketch.") - return None - reason = "no suffix of the executed sketch refines from here" - else: - reason = "no execution replans left" - remaining = list(exec_plan[failed_idx + 1:]) - logging.warning( - "Subgoal divergence after step %d (%s): %s. Resuming the " - "remaining %d step(s) open-loop; the divergence stands " - "recorded and the goal check decides the episode.", failed_idx, - failed_name, reason, len(remaining)) - return self._plan_to_policy(remaining, sketch=steps[failed_idx + 1:]) - - def _nudge_final_submission(self) -> Optional[Callable[[State], Action]]: - """One short follow-up query on the LAST attempt, after its query ended - with no captured plan: tell the agent to submit its best plan now. - - A session that hits the turn cap mid-iteration contributes - nothing, even when it has a near-working plan in context; this - converts that dead end into a submission attempt at the cost of a - few turns. - - The submitted plan is captured and executed even if its belief - rollout does not reach the goal, is scored a non-solve by the - task evaluator, or is flaky: there are no restarts left, and a - partial plan beats forfeiting the task. - """ - nudge = (_FINAL_SUBMIT_POLICY_NUDGE - if CFG.agent_solve_policy_mode else _FINAL_SUBMIT_NUDGE) - if CFG.agent_solve_use_journal: - nudge += (" If an earlier attempt's entry in the Attempt Log " - "records a better plan (captured or refused) than " - "anything from this attempt, resubmit that plan instead." - " After the submission, append ONE short factual entry " - "to ./journal.md for later fresh-context attempts and " - "tasks: what you tried (exact parameters), the key " - "measurements, and what to try differently - facts and " - "measurements only, no verdicts like 'impossible'.") - # SUSPEND (not clear) the attempt deadline for the nudge query: - # its cooperative refusals and the sandbox interrupt backstop - # must not block the submission (or the journal entry) itself. - # Restore it afterwards rather than leaking the None: _solve's - # per-attempt bookkeeping owns clearing the deadline, and a - # helper that silently disarms the wall clock is a trap for any - # future caller that runs mid-attempt. - saved_deadline = self._tool_context.attempt_deadline - self._tool_context.attempt_deadline = None - self._tool_context.capture_best_effort_plan = True - try: - nudge_responses = self._query_agent_sync(nudge, kind="test") - dead = query_fatal_error(nudge_responses) - if dead is not None: - raise AgentSessionFatalError( - "final-submission nudge died without the agent doing " - f"any work ({dead}); not recording this attempt as a " - "failure.") - except AgentSessionFatalError: - raise - except Exception as e: # pylint: disable=broad-except - logging.warning("Final-submission nudge failed: %s", e) - finally: - self._tool_context.capture_best_effort_plan = False - self._tool_context.attempt_deadline = saved_deadline - policy = self._consume_validated_plan() - if policy is not None: - logging.info( - "[%s] Final-submission nudge produced a validated plan.", - self._run_id) - return policy - - def _consume_validated_plan(self) -> Optional[Callable[[State], Action]]: - """Return a policy from an agent-validated plan, or None. - - ``submit_plan`` records a captured (goal-reaching, validated) - plan on the current solve task into the tool context. Returning - that exact simulator-verified plan guarantees the agent's tool- - validated answer is what executes, and avoids a fresh refinement - that with a different seed might not reproduce it. - """ - capture = self._tool_context.take_plan_capture() - if capture.policy_source: - return self._policy_capture_to_policy(capture) - if not capture.plan: - return None - # A capture with reached_goal False was accepted under the - # best-effort nudge; anything else is a validated solve. - validated = capture.reached_goal is not False - lines = list( - bilevel_sketch.format_plan_lines(capture.plan, - sketch=capture.sketch)) - self._last_capture_info = _CaptureInfo( - validated=validated, - reward=capture.eval_reward, - plan_lines=lines, - validation_summary=capture.validation_summary) - verdict = ("simulator-verified" if validated else - "best-effort: not a validated solve in the belief rollout") - # Log the full plan (options + continuous params + subgoal - # annotations) so the run log shows exactly what will execute. - logging.info( - "[%s] Using agent-validated plan from capture " - "(%d steps, %s):\n%s", self._run_id, len(capture.plan), verdict, - "\n".join(lines)) - return self._plan_to_policy(capture.plan, sketch=capture.sketch) - - def _policy_capture_to_policy(self, - capture: Any) -> Callable[[State], Action]: - """Turn a captured policy.py source into the execution policy. - - Policy-mode counterpart of the plan branch below: records the - capture metadata (the journal gets the source hash + validation - summary rather than plan lines), composes the SNAPSHOTTED source - against the real task and vocabulary, and wraps it in the - closed-loop executor. No ``SubgoalExecutionStatus`` is ever - published: the monitor and the divergence-replan path stay inert - (the constructor enforces replans == 0 in policy mode). - """ - # pylint: disable-next=import-outside-toplevel - from predicators.agent_sdk.policy_execution import \ - build_policy_option_fn - source = capture.policy_source - validated = capture.reached_goal is not False - sha = hashlib.sha256(source.encode("utf-8")).hexdigest()[:12] - n_lines = len(source.splitlines()) - self._last_capture_info = _CaptureInfo( - validated=validated, - reward=capture.eval_reward, - plan_lines=[ - f"" - ], - validation_summary=capture.validation_summary) - verdict = ("simulator-verified" if validated else - "best-effort: not a validated solve in the belief rollout") - logging.info( - "[%s] Using agent-validated POLICY from capture (sha=%s, " - "%d lines, %s).", self._run_id, sha, n_lines, verdict) - task = self._tool_context.current_task - assert task is not None - option_fn, err = build_policy_option_fn( - source, - task, - predicates=self._get_all_predicates(), - options=self._get_all_options(), - types=self._types) - if err is not None or option_fn is None: - raise ApproachFailure( - f"Captured policy.py failed to load for execution: {err}") - return self._policy_to_execution_policy(option_fn) - - def _policy_to_execution_policy( - self, option_fn: Any) -> Callable[[State], Action]: - """Closed-loop real executor for a composed policy option fn. - - Mirrors ``execute_policy_forward``'s semantics on the real env: - option execution failures (non-initiable, a skill raising - mid-execution - e.g. a motion-planning refusal - or an option - step-cap timeout) do NOT end the episode; the failure text is - surfaced to the policy via ``memory['last_failure']`` and the - next option is requested from the current state, bounded by - ``CFG.agent_policy_max_options`` total options, - ``CFG.agent_policy_max_repeated_failures`` consecutive failures - of one identical command (the stuck-loop guard), and - ``CFG.agent_policy_max_repeated_noops`` consecutive clean - completions of one identical command that changed nothing - observable (the guard's livelock twin). ``get_option`` - bugs and DONE end the episode via ``ApproachFailure`` (harmless - when the goal already holds). - - Implementation note: each issued option still runs through - ``utils.option_policy_to_policy`` (per-option step caps and the - Wait atom-change machinery live there), but the wrapper is - REBUILT after every surfaced failure - the failed option is - stuck inside the old wrapper's closure (its terminal never - holds), so a fresh wrapper is what makes the next call request a - new option. - """ - # pylint: disable-next=import-outside-toplevel - from predicators.agent_sdk.policy_execution import PolicyError, \ - option_repeat_key, repeated_failure_message, \ - repeated_noop_message, states_features_allclose - predicates = self._get_all_predicates() - - def _abstract(s: State) -> Set[GroundAtom]: - return utils.abstract(s, predicates) - - issued = 0 - last_failure: Optional[str] = None - repeat_key: Optional[Any] = None - repeat_count = 0 - issued_option: Optional[_Option] = None - issued_state: Optional[State] = None - noop_key: Optional[Any] = None - noop_count = 0 - - class _PolicyFatal(utils.OptionExecutionFailure): - """DONE / policy bug / budget: never surfaced, ends episode.""" - - def _option_policy(state: State) -> _Option: - nonlocal issued, last_failure, repeat_key, repeat_count, \ - issued_option, issued_state, noop_key, noop_count - if last_failure is None and issued > 0: - # The previous option completed cleanly: the policy is - # making progress, so the stuck-loop counter resets. - repeat_key = None - repeat_count = 0 - # ...unless the clean completion changed nothing - # observable: an identical command re-completing as a - # no-op K times is a livelock the failure guard cannot - # see (mirrors execute_policy_forward). - if issued_option is not None and issued_state is not None \ - and states_features_allclose(issued_state, state): - key = option_repeat_key(issued_option) - noop_count = noop_count + 1 if key == noop_key else 1 - noop_key = key - if noop_count >= CFG.agent_policy_max_repeated_noops: - raise _PolicyFatal( - repeated_noop_message(issued_option, noop_count)) - else: - noop_key = None - noop_count = 0 - if issued >= CFG.agent_policy_max_options: - raise _PolicyFatal( - "Policy exhausted its option budget " - f"({CFG.agent_policy_max_options} options) without " - "signalling DONE.") - try: - nxt = option_fn(state, last_failure) - except PolicyError as e: - raise _PolicyFatal(f"policy error: {e}") from e - if nxt is None: - logging.info("Policy signaled DONE after %d options.", issued) - raise _PolicyFatal("Policy signaled DONE.") - issued += 1 - last_failure = None - logging.info("Executing policy option %d/%d: %s", issued, - CFG.agent_policy_max_options, nxt.simple_str()) - if not nxt.initiable(state): - # Same text and attribution as execute_policy_forward - # (option_policy_to_policy's own "Unsound option policy" - # raise would name the PREVIOUS option in its info). - raise utils.OptionExecutionFailure( - "not initiable", info={"last_failed_option": nxt}) - issued_option = nxt - issued_state = state - return nxt - - def _fresh_inner() -> Callable[[State], Action]: - return utils.option_policy_to_policy( - _option_policy, - max_option_steps=CFG.max_num_steps_option_rollout, - abstract_function=_abstract) - - inner_box = {"inner": _fresh_inner()} - - def _execution_policy(state: State) -> Action: - nonlocal last_failure, repeat_key, repeat_count - while True: - try: - return inner_box["inner"](state) - except _PolicyFatal as e: - raise ApproachFailure(str(e)) from e - except utils.OptionExecutionFailure as e: - # Surface to the policy and continue: the failed - # option is stuck in the old wrapper, so rebuild. - failed = getattr(e, "info", {}).get("last_failed_option") - prefix = (f"{failed.name}: " if failed is not None else "") - last_failure = f"{prefix}{e}" - logging.info("Option failure surfaced to the policy: %s", - last_failure) - # Mirrors execute_policy_forward: K consecutive - # failures of one identical command are a policy - # bug, not recovery - end the episode attributably - # instead of burning the remaining option budget. - if failed is not None: - key = option_repeat_key(failed) - repeat_count = (repeat_count + - 1 if key == repeat_key else 1) - repeat_key = key - if (repeat_count >= - CFG.agent_policy_max_repeated_failures): - raise ApproachFailure( - repeated_failure_message(failed, - repeat_count)) from e - else: - repeat_key = None - repeat_count = 0 - inner_box["inner"] = _fresh_inner() - - return _execution_policy - - def _plan_to_policy( - self, - plan: List[_Option], - sketch: Optional[List[_SketchStep]] = None, - ) -> Callable[[State], Action]: - """Wrap a grounded option plan into a step-by-step policy. - - With ``CFG.agent_bilevel_max_execution_replans > 0`` and a full - per-step sketch, the policy also publishes a live - ``SubgoalExecutionStatus`` (via - ``get_execution_monitoring_info``) that the subgoal_annotations - execution monitor reads to check, at each option boundary, that - the just-finished step's annotation holds in the REAL state. On - divergence the monitor makes CogMan re-invoke solve(), which - lands in :meth:`_maybe_replan_from_divergence`. - """ - predicates = self._get_all_predicates() - - def _abstract(s: State) -> Set[GroundAtom]: - return utils.abstract(s, predicates) - - monitored = (CFG.agent_bilevel_max_execution_replans > 0 - and sketch is not None and len(sketch) == len(plan)) - - queue = list(plan) - total = len(queue) - status: Optional[SubgoalExecutionStatus] = None - if monitored: - assert sketch is not None - status = SubgoalExecutionStatus(sketch=list(sketch)) - self._exec_status = status - self._exec_plan = list(plan) - - def _option_policy(state: State) -> _Option: - del state # unused - if not queue: - logging.info("Option plan exhausted after %d options.", total) - # See the twin of this raise in utils.option_plan_to_policy: - # a finished plan and a failed option arrive as the same - # exception type, so the normal terminus is flagged. - raise utils.OptionExecutionFailure( - "Option plan exhausted!", info={"plan_exhausted": True}) - option = queue.pop(0) - num_done = total - len(queue) - if status is not None: - status.steps_initiated = num_done - status.current_option = option - next_option = None if not queue else queue[0].simple_str() - logging.info("Executing option %d/%d: %s (remaining=%d, next=%s)", - num_done, total, option.simple_str(), len(queue), - next_option) - return option - - inner = utils.option_policy_to_policy( - _option_policy, - max_option_steps=CFG.max_num_steps_option_rollout, - abstract_function=_abstract) - return self._wrap_option_failures(inner) - - def _replan_suffix( - self, - state: State, - task: Task, - sketch: List[_SketchStep], - failed_idx: int, - timeout: int, - ) -> Optional[Callable[[State], Action]]: - """Cheap-first recovery: re-refine a suffix of the current sketch. - - Divergence is usually a continuous-execution problem (a sampled - parameter whose real outcome differed from the option-model - rollout), not a wrong skeleton, so before paying for a fresh - agent sketch we retry the one we have. Candidate resume points - run from the failed step backward to just after the latest - earlier annotated step whose subgoals still hold in the current - state. The holds-check only bounds the walk-back - annotations - are optional and can hold coincidentally (e.g. a final - SwitchOff's {Off} atom holds before the switch was ever touched) - - so every candidate suffix must still refine AND forward- - validate from the current state before we trust it. Returns None - when no suffix candidate validates. - """ - assert self._option_model is not None - sub_task = Task(state, task.goal) - resume_floor = 0 - for j in range(failed_idx - 1, -1, -1): - step = sketch[j] - if step.subgoal_atoms is None and step.subgoal_neg_atoms is None: - continue - pos_ok = all(a.holds(state) for a in (step.subgoal_atoms or set())) - neg_ok = not any( - a.holds(state) for a in (step.subgoal_neg_atoms or set())) - if pos_ok and neg_ok: - resume_floor = j + 1 - break - start = time.perf_counter() - for j in range(failed_idx, resume_floor - 1, -1): - remaining = timeout - (time.perf_counter() - start) - if remaining <= 0: - break - suffix = list(sketch[j:]) - plan, success = self._refine_sketch(sub_task, - suffix, - remaining, - attempt=j) - if not success: - logging.info( - "Suffix replan: refinement failed resuming at " - "step %d.", j) - continue - ok, reason = bilevel_sketch.validate_plan_forward( - sub_task, - plan, - self._option_model, - predicates=self._get_all_predicates(), - sketch=suffix, - run_id=self._run_id, - ) - if ok: - logging.info( - "Suffix replan: resuming executed sketch at step %d " - "(%d steps).", j, len(plan)) - return self._plan_to_policy(plan, sketch=suffix) - logging.info( - "Suffix replan: forward validation failed resuming at " - "step %d: %s", j, reason) - return None diff --git a/predicators/approaches/agent_model_free_approach.py b/predicators/approaches/agent_model_free_approach.py index 90e9c43e63..9c9b08913e 100644 --- a/predicators/approaches/agent_model_free_approach.py +++ b/predicators/approaches/agent_model_free_approach.py @@ -1,63 +1,32 @@ -"""Agent model-free approach: fixed-vocabulary open-loop planning. +"""The base of the agent arms: the agent session, its tool context, the +recorded trajectories and the checkpoints. -The agent plans directly from its own world knowledge - no simulator -(model) is required to validate a plan before execution. Combines online -trajectory collection (via AgentModelFreeExplorer) with open-loop option plan -generation (via Claude Agent SDK). No predicate/process/type invention - -just stores trajectories and generates plans. - -Registered under the CLI approach name ``agent_model_free`` -(``agent_planner`` is kept as a deprecated alias). - -Example command: - python predicators/main.py --env pybullet_domino \ - --approach agent_model_free --seed 0 \ - --num_train_tasks 1 --num_test_tasks 1 \ - --num_online_learning_cycles 1 --explorer agent_model_free +The arms themselves (``agent_continual_approach`` and its siblings) add +the play loop of the continual protocol through ``ContinualPlayMixin``; +this class and the simulator-learning classes between it and the arms +carry the machinery they share. None of them solves a task on its own. """ -import copy import datetime import logging import os -import subprocess -from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, \ - cast +from typing import Any, Callable, Dict, List, Optional, Set, cast import dill as pkl -import numpy as np from gym.spaces import Box from predicators import utils -from predicators.agent_sdk.rendering import save_task_state_image -from predicators.agent_sdk.response_parser import extract_final_text -from predicators.agent_sdk.session_base import AgentSessionFatalError, \ - query_fatal_error -from predicators.agent_sdk.sketch_prompts import summarize_trajectories -from predicators.agent_sdk.tools import agent_render_resolution -from predicators.agent_sdk.tools.digests import render_options_digest, \ - render_types_digest from predicators.approaches import ApproachFailure from predicators.approaches.agent_session_mixin import AgentSessionMixin from predicators.approaches.base_approach import BaseApproach -from predicators.explorers import create_explorer -from predicators.explorers.base_explorer import BaseExplorer -from predicators.ground_truth_models import \ - augment_state_with_helper_objects, augment_task_with_helper_objects, \ - merge_gt_helper_predicates, merge_gt_helper_types from predicators.option_model import _OptionModelBase, create_option_model from predicators.settings import CFG -from predicators.structs import Action, Dataset, GroundAtom, \ - InteractionRequest, InteractionResult, LowLevelTrajectory, Object, \ - ParameterizedOption, ParameterizedSampler, Predicate, State, Task, Type +from predicators.structs import Action, Dataset, LowLevelTrajectory, \ + ParameterizedOption, Predicate, State, Task, Type class AgentModelFreeApproach(AgentSessionMixin, BaseApproach): - """Fixed-vocabulary open-loop planning via Claude Agent SDK. - - - Collects trajectories online using AgentModelFreeExplorer - - At solve time, queries the agent for an option plan - - No predicate/process/type invention - """ + """The agent session, tool context, trajectories and checkpoints the agent + arms share.""" def __init__(self, initial_predicates: Set[Predicate], @@ -70,16 +39,6 @@ def __init__(self, **kwargs: Any) -> None: super().__init__(initial_predicates, initial_options, types, action_space, train_tasks, *args, **kwargs) - # Optionally hand the agent the ground-truth helper scaffolding (e.g. - # the domino/fan grid loc/side types and grid predicates) so an - # "agent-with-grid" ablation plans over the oracle's vocabulary. Opt-in - # via CFG.use_gt_helpers; no-op for envs without a helper factory or - # when off. Merge here (before the agent session inits below) so the - # session, solve-time abstraction, and _get_all_predicates see them. - if self._use_gt_helpers(): - self._types = merge_gt_helper_types(self._types, CFG.env) - self._initial_predicates = merge_gt_helper_predicates( - self._initial_predicates, CFG.env) self._offline_dataset = Dataset([]) self._online_trajectories: List[LowLevelTrajectory] = [] self._option_model: Optional[_OptionModelBase] = ( @@ -87,65 +46,24 @@ def __init__(self, self._create_planner_option_model()) # Terminate Wait on atom change using the approach's predicates (which # may include invented ones), looked up lazily so the lambda picks up - # predicates invented after __init__. When the grid ablation is on, - # re-derive helper objects first so grid predicates (e.g. BallAtLoc) - # stay evaluable on the otherwise helper-free execution states. + # predicates invented after __init__. if self._option_model is not None and \ CFG.wait_option_terminate_on_atom_change: cast( # pylint: disable=protected-access Any, self._option_model)._abstract_function = ( - lambda s: utils.abstract(self._maybe_augment_state(s), - self._get_all_predicates())) - self._online_learning_cycle = 0 - # Synthesized per-skill samplers (option name -> sampler). Empty for - # the base planner; learning subclasses populate it. Threaded into - # bilevel refinement via _get_all_samplers() so continuous-parameter - # search aims at each step's subgoal instead of drawing uniformly. - self._synthesized_samplers: Dict[str, ParameterizedSampler] = {} - self._requests_train_task_idxs: Optional[List[int]] = None + lambda s: utils.abstract(s, self._get_all_predicates())) self._run_id = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - self._pre_test_conversation_log: Optional[List[Dict[str, Any]]] = None - # True only between begin_test_phase / end_test_phase, so per-episode - # hooks can act on test solves without touching exploration episodes. - self._in_test_phase = False - # 0-based index of the test task being solved, mirroring main.py's - # ``test_task_idx``. Incremented per test solve; threaded into the - # session-log filename via the ToolContext. - self._test_task_idx = -1 - # Solve-journal snapshot taken at begin_test_phase (None = no - # journal file existed) and whether it was captured successfully. - # end_test_phase archives the full test-phase journal outside the - # sandbox and rolls the file back to this snapshot, so learning - # entries persist across cycles while one evaluation's test-task - # entries never leak into the next evaluation. - self._pre_test_journal: Optional[str] = None - # Sandbox commit taken at begin_test_phase; end_test_phase - # archives what the phase added and resets the tree to it. - self._pre_test_sandbox_rev: Optional[str] = None - self._pre_test_attempts: Optional[str] = None - self._pre_test_journal_valid = False - # Scene renders attempted this episode. The first is the true initial - # state; later ones come from mid-episode replans and get distinct - # filenames so they don't overwrite the init snapshot. Reset in - # reset_for_new_episode. - self._episode_scene_renders = 0 - # Filename of the most recently saved scene render, consumed by - # _initial_image_section so the prompt references the image matching - # the state the agent is actually planning from. - self._last_scene_image_name: Optional[str] = None - # Initializes _tool_context and _agent_session_id (see mixin). Use the - # (possibly helper-augmented) vocabulary so the agent session exposes - # the grid types/predicates when CFG.use_gt_helpers is on. + # Initializes _tool_context and _agent_session_id (see mixin). self._init_agent_session_state(self._types, self._initial_predicates, initial_options, train_tasks) # Capture the underlying env once at construction. The initial option # model wraps ``env.simulate`` (a bound method), so ``__self__`` is the - # env. Later cycles may rebuild ``_option_model`` with a plain learned - # simulator that has no ``__self__``; pinning the env reference here - # keeps scene rendering (the probe's sim.render) working in every - # synthesis/solve cycle. + # env. A later model may rebuild ``_option_model`` with a plain + # learned simulator that has no ``__self__``; pinning the env + # reference here keeps scene rendering (the probe's sim.render) + # working in every round. env_self = getattr(getattr(self._option_model, '_simulator', None), '__self__', None) if env_self is not None: @@ -170,29 +88,6 @@ def _get_log_dir(self) -> str: # Overridable helpers (for subclass customisation) # ------------------------------------------------------------------ # - def _use_gt_helpers(self) -> bool: - """Whether to hand the agent the ground-truth helper scaffolding. - - Opt-in via ``CFG.use_gt_helpers`` (the process-planning - approaches read it too). When on, the grid helper - types/predicates are merged into the agent's vocabulary and the - solved task is augmented with the grid objects + oracle goal - (see ``__init__`` / ``_solve``). - """ - return CFG.use_gt_helpers - - def _maybe_augment_state(self, state: State) -> State: - """Re-derive GT helper objects on a state when the ablation is on. - - Executed states are helper-free (the grid is injected only into - the planning task), so this keeps helper predicates evaluable - during execution and Wait-on-atom-change termination. No-op when - helpers are disabled or the env has no helper factory. - """ - if self._use_gt_helpers(): - return augment_state_with_helper_objects(state, CFG.env) - return state - def _get_all_options(self) -> Set[ParameterizedOption]: """Return the full set of options available for planning.""" return self._initial_options @@ -201,28 +96,17 @@ def _get_all_predicates(self) -> Set[Predicate]: """Return the full set of predicates for abstraction.""" return self._initial_predicates - def _get_all_samplers(self) -> Dict[str, ParameterizedSampler]: - """Return synthesized per-skill samplers (option name -> sampler). - - Empty by default; learning subclasses populate the backing - field. Threaded into bilevel refinement so parameter search aims - at each step's subgoal. - """ - return self._synthesized_samplers - def _get_all_trajectories(self) -> List[LowLevelTrajectory]: """Return all trajectories (offline + online).""" return self._offline_dataset.trajectories + self._online_trajectories def _create_planner_option_model(self) -> Optional[_OptionModelBase]: - """Build the option model the planner tests plans against. + """Build the option model the tools roll plans out through. Honors two CFG knobs: * ``agent_planner_use_simulator`` -- when False, returns ``None`` - so the agent gets no ``submit_plan`` rollouts and must - plan open-loop from data + LLM reasoning (the model-free - baseline). + so the agent has no simulator to roll plans out in. * ``agent_planner_use_base_simulator`` -- when True (and a simulator is used), wraps the *base* env (``skip_residual_dynamics=True``), denying the planner the delayed @@ -234,875 +118,24 @@ def _create_planner_option_model(self) -> Optional[_OptionModelBase]: CFG.option_model_name, skip_residual_dynamics=CFG.agent_planner_use_base_simulator) - # ------------------------------------------------------------------ # - # AgentSessionMixin hooks - # ------------------------------------------------------------------ # - - # -- Prompt building blocks ----------------------------------------- # - - _SYSTEM_PROMPT_BASE = ( - "You are a planning agent. You observe task environments through " - "inspection tools and generate option plans to achieve goals. " - "You have access to read-only tools to inspect predicates, " - "options, trajectories, and training tasks. Use these to " - "understand the environment and generate effective plans.\n\n" - "Some effects may not be immediate - if an action triggers a " - "delayed process (e.g. water filling, dominoes cascading, " - "heating), insert a Wait after it so the effect has time " - "to occur before the next action. The Wait action holds the " - "robot's current pose. You can annotate Wait with target atoms " - "using `-> {atoms}` to specify exactly when it should terminate " - "(e.g. `Wait(robot:Robot)[] -> {Boiled(water:water_type)}`). " - "Use `NOT Pred(...)` for atoms that should become false. " - "If no annotation is provided, the Wait terminates on any atom " - "change. Without a Wait, the robot will proceed to the next " - "action before the delayed effect has occurred, which might " - "cause the plan to fail.") - - _SCRATCHPAD_SECTION = """ -## Scratchpad - CRITICAL -You MUST maintain `./notes.md` as your working memory. \ -**Read it at the very start of the session** and **read it \ -again before every submit_plan call** to remind yourself \ -what you already tried. **Update it immediately after every \ -submit_plan call** - no exceptions. - -Use this exact format for each option you are tuning: - -``` -## - Parameter Search -| # | params | outcome | notes | -|---|--------|---------|-------| -| 1 | [x, y, ...] | IK fail | ... | -| 2 | [x, y, ...] | success, JugNotAt... | ... | -``` - -After every test, append a row and update these summary fields: -- **Confirmed working params**: (list any that achieve the desired atoms) -- **Explored ranges**: e.g. "x: 0.9–1.05, y: 1.4–1.55" - look for GAPS -- **Unreachable region**: e.g. "y > 1.47 always IK-fails" -- **Next hypothesis**: what to try and why - -The cycle is: Read notes → plan next experiment → run test → \ -update notes → repeat. Without this loop you WILL forget what \ -you tried and repeat the same failed parameters. Treat notes.md \ -as your lab notebook - write after every single experiment. - -**If you notice you have NOT updated notes after a test, STOP \ -and update before doing anything else.**""" - - # -- System prompt --------------------------------------------------- # - - def _get_agent_system_prompt(self) -> str: - use_scratchpad = CFG.agent_planner_use_scratchpad - - sections = [self._SYSTEM_PROMPT_BASE] - - # Scratchpad - if use_scratchpad: - sections.append(self._SCRATCHPAD_SECTION) - - # Tuning workflow (numbered steps, dynamic) - steps = [] - if use_scratchpad: - steps.append( - "**Read `./notes.md` before every test**, then **update it " - "immediately after every submit_plan call**. Record " - "what you tried, what happened, and what you learned. " - "This is your memory - without it you will repeat failures.") - steps += [ - "**Review past session logs** in `./session_logs/` if available. " - "Previous queries and tool results from earlier sessions are " - "saved there. Read them to build on prior knowledge.", - "**Inspect rendered images** from `./test_images/` when " - "something goes wrong to understand the actual outcome.", - "**Expect geometric offsets.** The target position for " - "options is often offset from the reference object's reported " - "position due to object geometry. Explore a wide range around " - "the object's coordinates, not just values close to the " - "reported position.", - "**Search coarse-to-fine.** For each continuous parameter, " - "start with a WIDE grid spanning most of the valid range " - "(e.g. test 4–5 spread-out values across [low, high]). " - "Identify which coarse region works, THEN refine within it. " - "Never spend more than 3 attempts tweaking values in a small " - "neighborhood - if none work, jump to a different region. " - "Check your notes for gaps in the explored range.", - "**Vary ALL params, not just position.** Orientation and " - "other parameters change offsets and feasibility. If an " - "action fails at a position, try different values for the " - "other parameters before giving up on that region. Test at " - "least 2-3 values for each non-position parameter.", - ] - numbered = "\n".join(f"{i}. {s}" for i, s in enumerate(steps, 1)) - sections.append( - f"\n## Continuous Parameter Tuning\nFollow this workflow:\n" - f"{numbered}") - - return "\n".join(sections) - def _get_sandbox_reference_files(self) -> Dict[str, str]: """Document public control semantics without exporting implementation.""" return {"skills.md": "predicators/agent_sdk/prompts/public_skills.md"} - def _get_solve_tool_names(self) -> Optional[List[str]]: - # Type / option digests are static per session, so the solve - # prompt injects them directly (see _build_solve_prompt); the - # trajectory and task digests live in run_python's namespace - # (`trajectories` / `describe_trajectory` / `sim.task()`). - # Every remaining tool needs a simulator: submit_plan - # rolls fully-specified plans out through the option model and - # run_python probes it, so a planner without a simulator - # gets neither. - tools = [] - if CFG.agent_planner_use_simulator: - tools.append("submit_plan") - # Closed-loop policy mode: the delivery gate for the - # agent-written policy.py (submit_plan stays as a - # probe but no longer captures). - if CFG.agent_solve_policy_mode: - tools.append("submit_policy") - tools.append("run_python") - return tools - - # ------------------------------------------------------------------ # - # Learning - # ------------------------------------------------------------------ # - - def learn_from_offline_dataset(self, dataset: Dataset) -> None: - self._offline_dataset = dataset - self._tool_context.offline_trajectories = dataset.trajectories - if dataset.trajectories: - self._tool_context.example_state = \ - dataset.trajectories[0].states[0] - # Post-offline checkpoint: main.py's --load_approach path loads - # cycle None before the online loop, which previously had no - # file to read for this approach family (save only ran at the - # end of each online cycle). Hook so subclasses that learn more - # afterwards checkpoint once, after their own learning. - self._checkpoint_after_offline_learning() - - def get_interaction_requests(self) -> List[InteractionRequest]: - # Explore sessions carry their own phase tag (see the mixin's - # ``_explore_phase``) so their system prompt logs separately - # from the solve and synthesis ones. - self._explore_phase = True - try: - explorer = self._create_explorer() - requests: List[InteractionRequest] = [] - self._requests_train_task_idxs = [] - # A cycle's requests are all generated before any executes, so - # the explorer shows each query the plans already scheduled this - # cycle and asks for a complementary one. Fresh list per cycle. - self._tool_context.cycle_scheduled_plans = [] - for _ in range(CFG.online_nsrt_learning_requests_per_cycle): - task_idx = self._rng.choice(len(self._train_tasks)) - # Clear so a planning explorer's verdict is read fresh per - # request; non-planning explorers leave it None (no verdict). - self._tool_context.last_mental_model_solved = None - policy, termination_function = \ - explorer.get_exploration_strategy(task_idx, CFG.timeout) - req = InteractionRequest( - train_task_idx=task_idx, - act_policy=policy, - query_policy=lambda s: None, - termination_function=termination_function, - mental_model_solved=self._tool_context. - last_mental_model_solved) - requests.append(req) - self._requests_train_task_idxs.append(task_idx) - return requests - finally: - self._explore_phase = False - - def restore_interaction_requests(self, train_task_idxs: List[int]) -> None: - # A resume that reuses the cycle's persisted episodes never calls - # get_interaction_requests, which is what pairs each result with - # its train task below (run_20260828_173451 asserted here after a - # preemption mid-learn). - self._requests_train_task_idxs = list(train_task_idxs) - - def learn_from_interaction_results( - self, results: Sequence[InteractionResult]) -> None: - assert self._requests_train_task_idxs is not None - # Subclasses (e.g. AgentSimLearningApproach) may track the snapshot - # tags of the simulator/predicates files in effect when the explorer - # generated these plans. Tag each new trajectory so the next - # learn-phase prompt can surface provenance. ``None`` for approaches - # that don't track versions. - sim_version: Optional[str] = getattr(self, - "_current_simulator_version", - None) - preds_version: Optional[str] = getattr(self, - "_current_predicates_version", - None) - samplers_version: Optional[str] = getattr(self, - "_current_samplers_version", - None) - for i, result in enumerate(results): - task_idx = self._requests_train_task_idxs[i] - traj = LowLevelTrajectory( - result.states, - result.actions, - _train_task_idx=task_idx, - _source_simulator_version=sim_version, - _source_predicates_version=preds_version, - _source_samplers_version=samplers_version, - _env_reward=result.episode_reward, - _env_terminated=result.episode_terminated, - ) - self._online_trajectories.append(traj) - - # Update tool context - self._sync_tool_context() - - logging.info( - "[Run %s] Cycle %s: collected %d trajectories, %d total online.", - self._run_id, self._online_learning_cycle, len(results), - len(self._online_trajectories)) - - # Hook (default: save now) so subclasses that learn more after - # this method can checkpoint ONCE, after their learning, instead - # of writing a pre-learn file under the same cycle name that a - # preemption between the two writes would leave looking complete. - self._checkpoint_after_interaction_results(self._online_learning_cycle) - self._online_learning_cycle += 1 - - # ------------------------------------------------------------------ # - # Solving - # ------------------------------------------------------------------ # - - @staticmethod - def _wrap_option_failures( - policy: Callable[[State], Action]) -> Callable[[State], Action]: - """Wrap a policy so OptionExecutionFailure surfaces as ApproachFailure. - - Bilevel planning and the base open-loop planner both build a - low-level policy from a grounded option plan; this adapter is - their single place to translate the harness's option-execution - exception into the ApproachFailure CogMan expects. - """ - - def _policy(s: State) -> Action: - try: - return policy(s) - except utils.OptionExecutionFailure as e: - raise ApproachFailure(e.args[0], e.info) - - return _policy - def _solve(self, task: Task, timeout: int) -> Callable[[State], Action]: - self._sync_tool_context() - # When enabled, plan over the oracle's grid-augmented task: inject the - # grid loc/side objects and rewrite the goal to the grid BallAtLoc so - # the agent sees the oracle's scaffolding. Augmentation preserves - # goal_nl. No-op otherwise. - if self._use_gt_helpers(): - task = augment_task_with_helper_objects(task, CFG.env) - self._tool_context.current_task = task - # Render the initial state so the agent can see the scene layout. - self._render_initial_state_image(task) - try: - option_plan = self._query_agent_for_option_plan(task) - except AgentSessionFatalError: - # An ApproachFailure would be absorbed per-task; the broken - # session backend must terminate the run instead. - raise - except Exception as e: - raise ApproachFailure(f"Agent failed to produce option plan: {e}") - - preds = self._get_all_predicates() - policy = utils.option_plan_to_policy( - option_plan, - max_option_steps=CFG.max_num_steps_option_rollout, - abstract_function=lambda s: utils.abstract( - self._maybe_augment_state(s), preds)) - - return self._wrap_option_failures(policy) - - def _render_initial_state_image(self, task: Task) -> Optional[str]: - """Render the state this solve starts from and save to the sandbox. - - The first render of an episode is the true initial state - (``task{N:03d}_initial_state.png``); later renders come from - mid-episode replans and are saved as - ``task{N:03d}_replan{K}_state.png`` so they don't overwrite the - init snapshot (the replan "task" is rooted at the current, - partially-executed state). - - Returns the saved image path, or None if rendering is unavailable. - """ - self._last_scene_image_name = None - env = self._tool_context.env - if env is None: - return None - try: - # The session/sandbox (and thus ``image_save_dir`` on the - # ToolContext) is created lazily on the first agent query. This - # render runs *before* that query in ``_solve``, so on the very - # first test task the dir would still be None and task0's image - # would be silently skipped; ensure the session (and dir) exist - # first. Inside the try so a session-creation hiccup leaves - # rendering best-effort rather than crashing the solve. - self._ensure_agent_session() - except Exception as e: # pylint: disable=broad-except - logging.warning("Failed to render initial state image: %s", e) - return None - save_dir = self._tool_context.image_save_dir - if save_dir is None: - return None - task_id = self._tool_context.test_task_idx - replan_idx = self._episode_scene_renders - # Count attempts, not successes: if the init render fails, a later - # replan render still must not masquerade as the init image. - self._episode_scene_renders += 1 - if task_id is not None: - stem = f"task{task_id:03d}" - else: - stem = "" - if replan_idx == 0: - filename = f"{stem}_initial_state.png" if stem \ - else "initial_state.png" - else: - filename = f"{stem}_replan{replan_idx}_state.png" if stem \ - else f"replan{replan_idx}_state.png" - with agent_render_resolution(): - saved_path = save_task_state_image(env, task, save_dir, filename) - if saved_path is not None: - self._last_scene_image_name = filename - return saved_path - - def _initial_image_section(self) -> str: - """Return a prompt section pointing at the current solve's rendered - scene image, or an empty string if none was rendered. - - ``_render_initial_state_image`` must have been called first; - this references whichever file that call saved (init or replan - snapshot), so replan queries point at the current scene rather - than the stale episode-init image. - """ - save_dir = self._tool_context.image_save_dir - img_name = self._last_scene_image_name - if not save_dir or img_name is None: - return "" - if not os.path.exists(os.path.join(save_dir, img_name)): - return "" - # cwd of the agent is the sandbox root, so reference test_images/. - return ("\n## Initial State Image\n" - "A rendering of the scene this plan starts from has been " - f"saved to `./test_images/{img_name}`. **Read this image " - "first** to understand the spatial layout before " - "planning.\n") - - # ------------------------------------------------------------------ # - # Test phase lifecycle - # ------------------------------------------------------------------ # - - def begin_test_phase(self) -> None: - """Snapshot the learning conversation log and solve journal.""" - self._in_test_phase = True - self._test_task_idx = -1 - if self._agent_session is not None: - self._pre_test_conversation_log = copy.deepcopy( - self._agent_session.conversation_log) - else: - self._pre_test_conversation_log = None - self._snapshot_journal_for_test_phase() - self._snapshot_sandbox_for_test_phase() - - def end_test_phase(self) -> None: - """Restore the conversation log and journal to pre-test state.""" - self._in_test_phase = False - self._tool_context.test_task_idx = None - if self._agent_session is not None \ - and self._pre_test_conversation_log is not None: - # In-place restore through the public property (it returns - # the live list), so any other holder of the reference sees - # the rollback too. - log = self._agent_session.conversation_log - log[:] = self._pre_test_conversation_log - self._pre_test_conversation_log = None - self._archive_and_rollback_test_journal() - self._archive_and_rollback_test_sandbox() - - def _eval_phase_label(self) -> str: - """Name of the evaluation phase now running, for archives. - - The 0-based cycle whose learning it evaluates (matching - main.py's "ONLINE LEARNING CYCLE i"): the counter has already - advanced past that cycle's learn, so subtract 1; the pre- - learning initial test is "initial". - """ - eval_cycle = self._online_learning_cycle - 1 - return "initial" if eval_cycle < 0 else f"cycle{eval_cycle}" - - def _snapshot_sandbox_for_test_phase(self) -> None: - """Commit the sandbox tree so end_test_phase can restore it.""" - # pylint: disable-next=import-outside-toplevel - from predicators.agent_sdk import sandbox_setup - self._pre_test_sandbox_rev = None - try: - self._pre_test_sandbox_rev = sandbox_setup.snapshot_sandbox( - self._tool_context.sandbox_dir) - except (OSError, subprocess.SubprocessError) as e: - logging.warning( - "[%s] Failed to snapshot the sandbox at test-phase start; " - "its test-phase files will NOT be rolled back: %s", - self._run_id, e) - - def _archive_and_rollback_test_sandbox(self) -> None: - """Archive everything the test phase wrote into the sandbox, then - restore the pre-test tree. - - The journal rollback above covers two files; this covers the - rest (test-phase session logs, scene images, notes, scripts and - plan files the agent wrote), so no evaluation leaves anything a - later learn/explore session or evaluation can read. The archive - lives in the run's log dir, outside the sandbox. - """ - rev = self._pre_test_sandbox_rev - self._pre_test_sandbox_rev = None - sandbox_dir = self._tool_context.sandbox_dir - if rev is None or not sandbox_dir: - return - # pylint: disable-next=import-outside-toplevel - from predicators.agent_sdk import sandbox_setup - archive_dir = os.path.join(self._get_log_dir(), - f"sandbox_eval_{self._eval_phase_label()}") - try: - archived = sandbox_setup.rollback_sandbox(sandbox_dir, rev, - archive_dir) - except (OSError, subprocess.SubprocessError) as e: - logging.warning( - "[%s] Failed to roll back the test-phase sandbox: %s", - self._run_id, e) - return - logging.info( - "[%s] Rolled the sandbox back to its pre-test snapshot; %d " - "test-phase file(s) archived to %s", self._run_id, len(archived), - archive_dir) - - def _journal_active(self) -> bool: - """Whether solve-journal entries can be written at all.""" - return bool(CFG.agent_solve_use_journal - and self._tool_context.sandbox_dir) - - def _snapshot_journal_for_test_phase(self) -> None: - """Capture the learning-only journal and attempt log at test start. - - The snapshots are what ``end_test_phase`` rolls both files back - to. A failed capture leaves ``_pre_test_journal_valid`` False so - the rollback is skipped rather than destroying learning entries. - """ - self._pre_test_journal = None - self._pre_test_attempts = None - self._pre_test_journal_valid = False - if not self._journal_active(): - return - # pylint: disable-next=import-outside-toplevel - from predicators.agent_sdk import journal as journal_mod - sandbox_dir = self._tool_context.sandbox_dir - try: - self._pre_test_journal = journal_mod.read_raw(sandbox_dir) - self._pre_test_attempts = journal_mod.read_raw( - sandbox_dir, filename=journal_mod.ATTEMPTS_FILENAME) - self._pre_test_journal_valid = True - except OSError as e: - logging.warning( - "[%s] Failed to snapshot the solve journal at test-phase " - "start; test-phase entries will NOT be rolled back: %s", - self._run_id, e) - - def _archive_and_rollback_test_journal(self) -> None: - """Archive the test-phase journal and attempt log, then roll back. - - Each evaluation must be independent of previous evaluations: - content written while solving test tasks (harness attempt-log - entries and the agent's own journal notes) would otherwise leak - this evaluation's test tasks into the next one. Learning content - - the pre-test snapshots - persists across cycles. Before the - rollback, both files (learning + this evaluation's additions) - are copied to the run's log dir, which lives outside the sandbox - so the agent cannot read them, for later inspection. - """ - if not self._pre_test_journal_valid: - return - snapshots = { - "journal": self._pre_test_journal, - "attempts": self._pre_test_attempts, - } - self._pre_test_journal = None - self._pre_test_attempts = None - self._pre_test_journal_valid = False - sandbox_dir = self._tool_context.sandbox_dir - if not self._journal_active(): - return - assert sandbox_dir is not None - # pylint: disable-next=import-outside-toplevel - from predicators.agent_sdk import journal as journal_mod - filenames = { - "journal": journal_mod.JOURNAL_FILENAME, - "attempts": journal_mod.ATTEMPTS_FILENAME, - } - # One archive per evaluation phase, named by the 0-based cycle - # whose learning it evaluates (matching main.py's "ONLINE - # LEARNING CYCLE i"). The counter has already advanced past that - # cycle's learn, so subtract 1; the pre-learning initial test - # archives as "initial". A same-cycle re-eval overwrites its own - # file. - label = self._eval_phase_label() - try: - for kind, filename in filenames.items(): - content = journal_mod.read_raw(sandbox_dir, filename=filename) - if content is not None: - archive_path = os.path.join(self._get_log_dir(), - f"{kind}_eval_{label}.md") - with open(archive_path, "w", encoding="utf-8") as f: - f.write(content) - logging.info("[%s] Archived the test-phase %s to %s", - self._run_id, filename, archive_path) - journal_mod.restore(sandbox_dir, - snapshots[kind], - filename=filename) - except OSError as e: - logging.warning( - "[%s] Failed to archive/roll back the test-phase solve " - "journal: %s", self._run_id, e) - - def reset_for_new_episode(self) -> None: - """Advance the test-task counter at each test episode start. - - CogMan calls this exactly once per test task (via - ``cogman.reset`` in main.py's ``_solve_task``) and never on a - replan inside an episode, so the counter stays in lockstep with - main.py's ``test_task_idx``. The index reaches the sandbox via - the ToolContext and lands in the session-log filename. No-op - outside the test phase. - """ - super().reset_for_new_episode() - # New episode -> the next scene render is a true init snapshot. - self._episode_scene_renders = 0 - if self._in_test_phase: - self._test_task_idx += 1 - self._tool_context.test_task_idx = self._test_task_idx - - def _query_agent_for_option_plan(self, task: Task) -> list: - """Query the agent for an option plan and parse it.""" - prompt = self._build_solve_prompt(task) - responses = self._query_agent_sync(prompt, kind="test") - dead = query_fatal_error(responses) - if dead is not None: - # An outage is not a failed attempt: recording 0/1 here - # would write a bogus eval datapoint. Stop the run; the - # relaunch re-runs this cycle's test. - raise AgentSessionFatalError( - "test query died without the agent doing any work " - f"({dead}); not recording this attempt as a failure.") - plan_text = self._extract_option_plan_text(responses) - - if not plan_text: - # Log the raw responses for debugging. - n_responses = len(responses) - types = [r.get("type") for r in responses] - raise ApproachFailure( - f"Agent returned empty plan text. " - f"Got {n_responses} responses with types: {types}") - - return self._parse_and_ground_plan(plan_text, task) - - def _solve_prompt_visualize_line(self) -> str: - """The stuck-step visualization bullet: the probe's staging + render is - the only visualization surface, so the bullet appears only when - run_python is offered.""" - if CFG.agent_planner_use_simulator: - return ( - "- **Use run_python when stuck** - after 3+ failures on " - "the same step, STOP testing and use run_python " - "(`sim.reset(mods={...})`, then `sim.render(...)`) to move " - "the object to several candidate positions and " - "orientations. It's free (no physics). Find the right " - "region visually, then test.\n") - return "" - - def _solve_prompt_scratchpad_line(self) -> str: - """Return the notes.md bullet for the solve prompt, or empty.""" - if CFG.agent_planner_use_scratchpad: - return ( - "- **Read `./notes.md` before every " - "submit_plan call** " - "and **update it immediately after each call** - append a " - "row to the parameter table and update the explored-ranges " - "summary. If you realize you forgot to update, STOP and " - "update before doing anything else.\n") - return "" - - def _build_solve_prompt(self, task: Task) -> str: - """Build the prompt for generating an option plan.""" - init_state = task.init - objects = list(init_state) - - # Objects - obj_strs = [] - for obj in sorted(objects, key=lambda o: o.name): - obj_strs.append(f" {obj.name}: {obj.type.name}") - - # Goal. Only expose goal atoms whose predicate is in the agent's - # current predicate set (same filter as the bilevel sketch - # prompt): approaches that strip env predicates rely on goal_nl - # to communicate the goal. - visible_preds = self._get_all_predicates() - goal_strs = [ - str(a) for a in sorted(task.goal, key=str) - if a.predicate in visible_preds - ] - - # Types and options: static per-session digests, injected here - # instead of costing a tool turn (see _get_solve_tool_names). - types_digest = render_types_digest(self._tool_context.types) - options_digest = render_options_digest( - self._get_all_options(), - gt_options_ref_path=self._tool_context.gt_options_ref_path) - - # Current atoms - atoms = utils.abstract(init_state, self._get_all_predicates()) - atom_strs = [str(a) for a in sorted(atoms, key=str)] - - # Trajectory summary - traj_summary = self._build_trajectory_summary() - - # State features (compact) - state_str = init_state.dict_str(indent=2) - - # Available tools - tool_names = self._get_solve_tool_names() - tools_str = "" - if tool_names: - tool_list = "\n".join(f" - {t}" for t in tool_names) - tools_str = f"\n## Available Tools\n{tool_list}\n" - - # Natural language goal description (if available) - goal_nl_section = "" - if task.goal_nl: - goal_nl_section = f""" -## Goal Description -{task.goal_nl} -""" - - # Initial state image reference - initial_image_section = self._initial_image_section() - - if CFG.agent_planner_use_simulator: - instructions_intro = ( - "Use your available tools to inspect the environment and " - "test your plan before committing to it.") - else: - instructions_intro = ( - "You do NOT have a simulator to test plans against. Inspect " - "the trajectory data and reason carefully about the dynamics, " - "then commit to your best open-loop plan.") - - prompt = f"""You are solving a task. \ -Generate an option plan to achieve the goal. -{goal_nl_section} -## Goal Atoms -{chr(10).join(goal_strs)} - -## Initial State Atoms -{chr(10).join(atom_strs)} - -## Initial State Features -{state_str} -{initial_image_section}{self._solve_prompt_extra_sections()} -## Objects -{chr(10).join(obj_strs)} - -## Object Types -{types_digest} - -## Available Options -{options_digest} -{traj_summary}{tools_str} -## Instructions -{instructions_intro} - -Based on the task information and any past trajectory data, output an option plan to achieve the goal. - -After any action whose desired subgoal depends on a delayed process (e.g. water \ -filling, dominoes cascading, heating), insert a Wait action to let the process \ -complete before proceeding. You can annotate Wait with target atoms using \ -`-> {{atoms}}` to specify exactly when it should terminate. Use `NOT Pred(...)` for \ -atoms that should become false. If no annotation is provided, the Wait terminates on \ -any atom change. Only use Wait when there is a genuine delayed effect; do not insert \ -it between actions with immediate effects (e.g. Pick, Place). - -For Wait with target atoms: `Wait(robot:Robot)[] -> {{Boiled(water:water_type)}}` -For negated targets: `Wait(robot:Robot)[] -> {{NOT Touching(a:block, b:block)}}` - -**Important - parameter tuning workflow:** -- When a step fails or produces unexpected results, inspect the rendered images \ -in `./test_images/` to see what actually happened in the scene. -{self._solve_prompt_scratchpad_line()}\ -- Review past session logs in `./session_logs/` if available - they contain prior queries and results. -- When a step fails (e.g. IK error), use the image + object poses to reason about \ -WHY and adjust params directionally. Don't just try random nearby values. -{self._solve_prompt_visualize_line()}\ -- **Vary all parameters, not just position** - orientation and other params affect \ -both the outcome and whether the action succeeds. Try 2-3 values for each \ -non-position parameter per target region. -- **Search coarse-to-fine**: spread initial attempts across the full parameter range. \ -If 3 nearby values all fail the same way, jump to a very different region instead of \ -continuing to tweak. Check your notes for gaps in explored ranges. - -Output the plan with one option per line in this exact format: - OptionName(obj1:type1, obj2:type2)[param1, param2] - -If an option has no continuous parameters, use empty brackets: OptionName(obj1:type1)[] - -Output ONLY the option plan lines at the end, after any analysis.""" - - return prompt - - def _build_trajectory_summary(self) -> str: - """Summarize trajectory data for context.""" - return summarize_trajectories(self._get_all_trajectories(), - self._get_all_predicates(), - train_tasks=self._train_tasks) - - def _solve_prompt_extra_sections(self) -> str: - """Subclass hook: sections inserted into the solve prompt after the - initial-state image reference (empty by default).""" - return "" - - @staticmethod - def _extract_option_plan_text(responses: List[Dict[str, Any]]) -> str: - """Extract plan text from the last assistant text response.""" - return extract_final_text(responses) - - @staticmethod - def _strip_code_fences(text: str) -> str: - """Strip markdown code fences wrapping the plan text.""" - lines = text.split('\n') - # Remove leading/trailing ``` lines (with optional language tag). - while lines and lines[0].strip().startswith('```'): - lines.pop(0) - while lines and lines[-1].strip().startswith('```'): - lines.pop() - return '\n'.join(lines) - - def _parse_wait_annotations( - self, - text: str, - predicates: Set[Predicate], - objects: Sequence[Object], - ) -> List[Tuple[Set[GroundAtom], Set[GroundAtom]]]: - """Parse ``-> {atoms}`` annotations from plan lines. - - Returns a list parallel to the option lines in the text. Each - entry is ``(positive_atoms, negative_atoms)`` for Wait lines - with annotations, or ``(set(), set())`` otherwise. - """ - results: List[Tuple[Set[GroundAtom], Set[GroundAtom]]] = [] - option_names = {o.name for o in self._get_all_options()} - for line in text.split('\n'): - stripped = line.strip() - if not stripped: - continue - first_token = stripped.split('(')[0] - if first_token not in option_names: - if results: - break - continue - if first_token == "Wait" and '->' in stripped: - pos, neg = utils.parse_wait_target_annotations( - stripped, predicates, objects) - results.append((pos, neg)) - else: - results.append((set(), set())) - return results - - def _parse_and_ground_plan(self, plan_text: str, task: Task) -> list: - """Parse option plan text and ground into executable options.""" - objects = list(task.init) - all_options = self._get_all_options() - option_names = sorted(o.name for o in all_options) - - # Strip markdown code fences that agents often wrap plans in. - cleaned_text = self._strip_code_fences(plan_text) - - # Extract Wait target annotations before stripping them. - wait_annotations = self._parse_wait_annotations( - cleaned_text, self._get_all_predicates(), objects) - - # Strip annotations so the option plan parser doesn't choke. - parseable_text = utils.strip_wait_annotations(cleaned_text) - - parsed = utils.parse_model_output_into_option_plan( - parseable_text, - objects, - self._types, - all_options, - parse_continuous_params=True) - if not parsed: - raise ApproachFailure(f"Parsed empty option plan from agent.\n" - f" Plan text:\n{plan_text}\n" - f" Available option names: {option_names}") - - grounded = [] - for i, (option, objs, params) in enumerate(parsed): - try: - params_arr = np.array(params, dtype=np.float32) - ground_opt = option.ground(objs, params_arr) - # Inject Wait target atoms from annotations. - if (ground_opt.name == "Wait" and i < len(wait_annotations)): - pos, neg = wait_annotations[i] - if pos: - ground_opt.memory["wait_target_atoms"] = pos - if neg: - ground_opt.memory["wait_target_neg_atoms"] = neg - grounded.append(ground_opt) - except Exception as e: # pylint: disable=broad-except - logging.warning("[Run %s] Failed to ground option " - "%s: %s", self._run_id, option.name, e) - break - - if not grounded: - raise ApproachFailure("No options successfully grounded.") - logging.info("[Run %s] Agent produced plan with %d options.", - self._run_id, len(grounded)) - return grounded - - # ------------------------------------------------------------------ # - # Explorer - # ------------------------------------------------------------------ # - - def _create_explorer(self) -> BaseExplorer: - """Create explorer for interaction requests.""" - if CFG.explorer in ("agent_model_free", "agent_model_based", - "agent_plan", "agent_bilevel"): - self._sync_tool_context() - return self._create_agent_explorer( - self._get_all_predicates(), - self._get_all_options(), - name=CFG.explorer, - ) - return create_explorer( - CFG.explorer, - self._get_all_predicates(), - self._get_all_options(), - self._types, - self._action_space, - self._train_tasks, - ) + raise ApproachFailure( + f"{self.get_name()} has no task solver: the agent arms play " + "levels under the continual protocol " + "(ContinualPlayMixin.play_level).") def _sync_tool_context(self) -> None: """Push current approach state into the shared ToolContext. - The MCP tools (submit_plan, run_python, etc.) read from the + The MCP tools (run_python and the play tools) read from the ToolContext dataclass, not the approach directly. This keeps them in sync after mutations (e.g. new trajectories collected, - options added). Called before each solve and learning - interaction. Subclasses should call super() and then set + options added). Subclasses should call super() and then set additional fields (e.g. skill_factory_context). """ self._tool_context.types = self._types @@ -1120,9 +153,6 @@ def _sync_tool_context(self) -> None: self._tool_context.online_trajectories = self._online_trajectories self._tool_context.log_dir = self._get_log_dir() self._tool_context.option_model = self._option_model - # Synthesized samplers, so the explorer and synthesis tools thread the - # same per-skill samplers into refinement that the approach uses. - self._tool_context.parameterized_samplers = self._get_all_samplers() # Wire the active-experiment info-gain scorer when a learning subclass # exposes one and info-seeking exploration is on. Syncing the bound # method (not a snapshot) keeps it pointed at the latest fit/ensemble. @@ -1174,24 +204,9 @@ def _load_extra_save_state(self, save_dict: Dict[str, Any]) -> None: been refreshed, but before the tool context is re-synced. """ - def _checkpoint_after_offline_learning(self) -> None: - """Checkpoint hook at the end of offline learning (see above).""" - self.save(None) - - def _checkpoint_after_interaction_results(self, cycle: int) -> None: - """Checkpoint hook at the end of an online cycle's data collection, - BEFORE the cycle counter increments (see above).""" - self.save(cycle) - def save(self, online_learning_cycle: Optional[int] = None) -> None: - """Save approach state to disk. - - The pickled ``online_learning_cycle`` is the cycle the FILE - denotes (``_c`` means "cycle c completed"), so a subclass that - saves after the counter already advanced still writes a - consistent checkpoint; ``_None`` (post-offline) records the live - counter (0). - """ + """Save approach state to disk; the continual runner names the + checkpoint after the level it closes (``online_learning_cycle``).""" save_path = utils.get_approach_save_path_str() path = f"{save_path}_{online_learning_cycle}.{self._save_suffix}" save_dict = { @@ -1199,9 +214,6 @@ def save(self, online_learning_cycle: Optional[int] = None) -> None: self._offline_dataset, "online_trajectories": self._online_trajectories, - "online_learning_cycle": - (online_learning_cycle if online_learning_cycle is not None else - self._online_learning_cycle), "run_id": self._run_id, "agent_session_id": @@ -1220,13 +232,6 @@ def load(self, online_learning_cycle: Optional[int] = None) -> None: self._offline_dataset = save_dict["offline_dataset"] self._online_trajectories = save_dict["online_trajectories"] - # ``_c`` means cycle c completed -> resume at c+1; the - # post-offline ``_None`` file means no online cycle completed -> - # resume at its recorded counter (0), NOT +1 (which would write - # cycle 0's checkpoint as ``_1`` and break the next load). - saved_cycle = save_dict["online_learning_cycle"] - self._online_learning_cycle = (saved_cycle + 1 if online_learning_cycle - is not None else saved_cycle) # pylint: disable=attribute-defined-outside-init # (_agent_session_id is initialized via the agent-session mixin.) self._agent_session_id = save_dict.get("agent_session_id") diff --git a/predicators/approaches/agent_nl_world_model_approach.py b/predicators/approaches/agent_nl_world_model_approach.py deleted file mode 100644 index 081b32ff14..0000000000 --- a/predicators/approaches/agent_nl_world_model_approach.py +++ /dev/null @@ -1,363 +0,0 @@ -"""The natural-language world model baseline (paper arm C3). - -The same loop and experiments as the code arms - the agent explores the -train tasks, learns after every cycle, and solves the test tasks - but -the learned model is a natural-language document, ``world_model.md``, -never executable code. The learn session writes it from the recorded -data with the same data tools the code arms get; the solve and explore -sessions receive it quoted into every task message and plan by -reasoning over it, with no simulator to test plans against (the -model-free planner's open-loop solve). Predicates are the env's kept -initial ones (the same allowlist as the code arms) and goals arrive as -natural language. -""" -from __future__ import annotations - -import logging -import os -from typing import Any, Dict, List, Optional, Sequence, Set - -import numpy as np - -from predicators.agent_sdk import learn_prompts -from predicators.agent_sdk.session_base import AgentSessionFatalError, \ - query_fatal_error -from predicators.agent_sdk.tools import _SnapshotTarget -from predicators.agent_sdk.tools.digests import render_options_digest, \ - render_trajectory_digest, render_types_digest -from predicators.agent_sdk.tools.python_exec import _make_python_exec_tool -from predicators.agent_sdk.tools.results import _make_coercing_tool, \ - _make_spilling_text_result -from predicators.agent_sdk.tools.snapshots import finalize_versioned_snapshot -from predicators.approaches.agent_model_free_approach import \ - AgentModelFreeApproach -from predicators.code_sim_learning.program_world_model import \ - option_transitions -from predicators.settings import CFG -from predicators.structs import Dataset, InteractionResult, \ - LowLevelTrajectory, Predicate - -logger = logging.getLogger(__name__) - -_NOTES_FILE = "world_model.md" -_NOTES_VERSIONS_DIR = "world_model_versions" - -_RUN_PYTHON_DESCRIPTION = ( - "Execute Python code (`code`, or `path` to a .py file you wrote in " - "the sandbox) for data exploration. Available variables: " - "trajectories (List[LowLevelTrajectory]; each has `is_demo`, " - "`train_task_idx`, `states`, `actions`; each action's `get_option()` " - "is the skill that produced it), train_tasks (List[Task]; each has " - "`init`, `goal`, `goal_holds(state)`), is_goal_state (callable: " - "state, task_idx -> bool), describe_trajectory(traj_idx, " - "include_states=True, include_atoms=False, max_timesteps=10), and " - "np. print() output is returned; the namespace persists across " - "calls; oversize output is saved under `tool_outputs/run_python/` " - "and previewed. There is no simulator in this session: the world " - "model you write is a document, checked by predicting recorded " - "transitions from it by hand.") - - -class AgentNotesWorldModelApproach(AgentModelFreeApproach): - """Model-free agentic planning over a learned natural-language world model - document.""" - - _save_suffix: str = "AgentNotesWM" - - def __init__(self, *args: Any, **kwargs: Any) -> None: - super().__init__(*args, **kwargs) - self._notes: str = "" - self._notes_version: Optional[str] = None - missing = [i for i, t in enumerate(self._train_tasks) if not t.goal_nl] - assert not missing, ( - f"{type(self).__name__} presents goals in natural language, so " - f"every train task must set `goal_nl`. Missing on task " - f"indices: {missing}") - - @classmethod - def get_name(cls) -> str: - return "agent_nl_world_model" - - # ── Vocabulary ─────────────────────────────────────────────── - - def _get_all_predicates(self) -> Set[Predicate]: - """The env predicates, restricted to the configured allowlist - (``agent_sim_learn_kept_predicates_names``) like the code arms.""" - preds = super()._get_all_predicates() - kept = CFG.agent_sim_learn_kept_predicates_names - if kept: - preds = {p for p in preds if p.name in set(kept)} - return preds - - # ── Session surface ────────────────────────────────────────── - - def _get_synthesis_tool_names(self) -> Optional[List[str]]: - return ["run_python"] - - def _get_agent_system_prompt(self) -> str: - if self._learning_mode: - return learn_prompts.build_notes_learn_system_prompt() - return "\n\n".join([ - super()._get_agent_system_prompt(), - learn_prompts.render_notes_solve_system_section(), - ]) - - def _solve_prompt_extra_sections(self) -> str: - return learn_prompts.render_world_model_notes_block( - self._notes, - self._notes_paths()["notes_file_for_agent"]) - - def _sync_tool_context(self) -> None: - super()._sync_tool_context() - self._tool_context.world_model_notes = self._notes - self._tool_context.world_model_notes_path = \ - self._notes_paths()["notes_file_for_agent"] - - # ── Learning ───────────────────────────────────────────────── - - def learn_from_offline_dataset(self, dataset: Dataset) -> None: - super().learn_from_offline_dataset(dataset) - self._learn_notes() - self.save(None) - - def learn_from_interaction_results( - self, results: Sequence[InteractionResult]) -> None: - cycle = self._online_learning_cycle - super().learn_from_interaction_results(results) - self._learn_notes() - self.save(cycle) - - def _checkpoint_after_offline_learning(self) -> None: - """No-op: this class checkpoints after its own learning.""" - - def _checkpoint_after_interaction_results(self, cycle: int) -> None: - """No-op: this class checkpoints after its own learning.""" - del cycle - - def _learn_notes(self) -> None: - """Run one document-writing session over all recorded data.""" - trajectories = self._get_all_trajectories() - if not trajectories and not CFG.agent_sim_learn_zero_shot: - logger.warning("No recorded trajectories; skipping the world " - "model document session.") - return - if not trajectories: - logger.info("Zero-shot synthesis: the agent writes the world " - "model document without data.") - self._run_notes_session(trajectories) - - def _notes_paths(self) -> Dict[str, str]: - """Host and agent-visible paths of the document (the residual arm's - sandbox mapping).""" - if CFG.agent_sdk_use_local_sandbox: - sandbox_dir: Optional[str] = os.path.abspath( - os.path.join(self._get_log_dir(), "sandbox")) - else: - sandbox_dir = self._tool_context.sandbox_dir - base = sandbox_dir or self._get_log_dir() - notes_file = os.path.join(base, _NOTES_FILE) - if CFG.agent_sdk_use_local_sandbox: - notes_file_for_agent = f"./{_NOTES_FILE}" - sandbox_dir_for_agent: Optional[str] = "." - elif sandbox_dir: - notes_file_for_agent = f"/sandbox/{_NOTES_FILE}" - sandbox_dir_for_agent = "/sandbox" - else: - notes_file_for_agent = notes_file - sandbox_dir_for_agent = None - return { - "base": base, - "notes_file": notes_file, - "versions_dir": os.path.join(base, _NOTES_VERSIONS_DIR), - "notes_file_for_agent": notes_file_for_agent, - "sandbox_dir_for_agent": sandbox_dir_for_agent or "", - } - - def _build_notes_exec_ns( - self, trajectories: List[LowLevelTrajectory]) -> Dict[str, Any]: - predicates = self._get_all_predicates() - train_tasks = self._train_tasks - - def describe_trajectory(traj_idx: int, - include_states: bool = True, - include_atoms: bool = False, - max_timesteps: int = 10) -> str: - return render_trajectory_digest(trajectories, - train_tasks, - predicates, - traj_idx, - include_states=include_states, - include_atoms=include_atoms, - max_timesteps=max_timesteps) - - return { - "trajectories": - trajectories, - "train_tasks": - train_tasks, - "is_goal_state": - lambda state, task_idx: train_tasks[task_idx].goal_holds(state), - "describe_trajectory": - describe_trajectory, - "np": - np, - } - - def _run_notes_session(self, - trajectories: List[LowLevelTrajectory]) -> None: - # pylint: disable=import-outside-toplevel - from claude_agent_sdk import tool as _sdk_tool - - from predicators.approaches.agent_sim_learning_approach import \ - AgentSimLearningApproach - - # pylint: enable=import-outside-toplevel - paths = self._notes_paths() - os.makedirs(paths["base"], exist_ok=True) - # A restored document is written back so the agent can Read it. - if self._notes and not os.path.isfile(paths["notes_file"]): - with open(paths["notes_file"], "w", encoding="utf-8") as f: - f.write(self._notes) - exec_ns = self._build_notes_exec_ns(trajectories) - run_python = _make_python_exec_tool( - _make_coercing_tool(_sdk_tool), - name="run_python", - description=_RUN_PYTHON_DESCRIPTION, - exec_ns=exec_ns, - sandbox_dir=paths["base"], - sandbox_dir_for_agent=paths["sandbox_dir_for_agent"] or None, - text_result=_make_spilling_text_result( - paths["base"], - agent_prefix=paths["sandbox_dir_for_agent"] or None), - call_timeout_s=CFG.agent_sdk_synthesis_python_call_timeout, - ) - ctx = self._tool_context - ctx.extra_mcp_tools = [run_python] - ctx.learn_cycle_index = self._online_learning_cycle - targets = [ - _SnapshotTarget( - live_file=paths["notes_file"], - versions_dir=paths["versions_dir"], - artifact_name="world_model_notes", - cycle_index_provider=lambda: self._online_learning_cycle, - ) - ] - build_hooks = ( - AgentSimLearningApproach._build_synthesis_session_hooks # pylint: disable=protected-access - ) - ctx.extra_session_hooks = build_hooks(targets, paths["base"]) - self._learning_mode = True - self._close_agent_session() - try: - self._ensure_agent_session() - message = self._build_notes_learn_message(trajectories, paths) - responses = self._query_agent_sync(message, kind="learn") - dead = query_fatal_error(responses) - if dead is not None: - raise AgentSessionFatalError( - "The learn session died without the agent doing any " - f"work ({dead}); refusing to checkpoint this cycle as " - "learned.") - finally: - ctx.extra_session_hooks = {} - ctx.extra_mcp_tools = [] - ctx.learn_cycle_index = None - self._learning_mode = False - self._close_agent_session() - self._load_notes(paths) - - def _build_notes_learn_message(self, - trajectories: List[LowLevelTrajectory], - paths: Dict[str, str]) -> str: - predicates = self._get_all_predicates() - n_trajs = len(trajectories) - n_demos = sum(1 for t in trajectories if t.is_demo) - n_transitions = sum( - len(option_transitions(t, predicates)) for t in trajectories) - session_tool_names = (self._agent_session.tool_names - if self._agent_session is not None else []) - extra_messages: List[str] = [] - if not trajectories and CFG.agent_sim_learn_zero_shot: - extra_messages.append( - learn_prompts.render_notes_zero_shot_message()) - listing = "\n".join( - f" [{i}] {'demo' if t.is_demo else 'interaction'}, task " - f"{t.train_task_idx}" for i, t in enumerate(trajectories)) - signatures = "\n".join( - f"- {p.name}({', '.join(t.name for t in p.types)})" - for p in sorted(predicates, key=lambda p: p.name)) - objective = next((t.evaluator.objective_description() - for t in self._train_tasks if t.evaluator is not None - and t.evaluator.objective_description()), "") - return learn_prompts.build_notes_learn_message( - n_trajs=n_trajs, - n_transitions=n_transitions, - n_demos=n_demos, - n_interaction=n_trajs - n_demos, - trajectory_listing=listing, - structs_ref=self._structs_reference_path(), - predicate_listing=signatures or "(none)", - types_digest=render_types_digest(self._tool_context.types), - options_digest=render_options_digest( - self._get_all_options(), - gt_options_ref_path=self._tool_context.gt_options_ref_path), - notes_file=paths["notes_file_for_agent"], - goal_nls=[t.goal_nl or "" for t in self._train_tasks], - has_prior_notes=os.path.isfile(paths["notes_file"]), - objective_block=learn_prompts.render_objective_block(objective), - tools_block=learn_prompts.render_tools_block(session_tool_names), - extra_messages=extra_messages, - ) - - def _structs_reference_path(self) -> str: - """Write the data-structures source into the sandbox reference dir (the - residual arm's convention) and return its agent-visible path.""" - # pylint: disable-next=import-outside-toplevel - import inspect - - # pylint: disable-next=import-outside-toplevel - from predicators import structs - paths = self._notes_paths() - ref_dir = os.path.join(paths["base"], "reference") - os.makedirs(ref_dir, exist_ok=True) - with open(os.path.join(ref_dir, "structs.py"), "w", - encoding="utf-8") as f: - f.write(inspect.getsource(structs)) - if paths["sandbox_dir_for_agent"]: - return f"{paths['sandbox_dir_for_agent']}/reference/structs.py" - return os.path.join(ref_dir, "structs.py") - - def _load_notes(self, paths: Dict[str, str]) -> None: - tag = finalize_versioned_snapshot( - paths["notes_file"], - paths["versions_dir"], - cycle_idx=self._online_learning_cycle, - artifact_name="world_model_notes") - if not os.path.isfile(paths["notes_file"]): - logger.warning( - "The session left no %s; the previous document " - "stands.", _NOTES_FILE) - return - with open(paths["notes_file"], "r", encoding="utf-8") as f: - self._notes = f.read() - self._notes_version = tag - self._sync_tool_context() - logger.info("Loaded the world model document (%d chars, %s).", - len(self._notes), tag or "unversioned") - - # ── Checkpointing ──────────────────────────────────────────── - - def _extra_save_state(self) -> Dict[str, Any]: - return { - "world_model_notes": self._notes, - "world_model_notes_version": self._notes_version, - } - - def _load_extra_save_state(self, save_dict: Dict[str, Any]) -> None: - self._notes = str(save_dict.get("world_model_notes") or "") - self._notes_version = save_dict.get("world_model_notes_version") - if self._notes: - paths = self._notes_paths() - os.makedirs(paths["base"], exist_ok=True) - with open(paths["notes_file"], "w", encoding="utf-8") as f: - f.write(self._notes) diff --git a/predicators/approaches/agent_program_world_model_approach.py b/predicators/approaches/agent_program_world_model_approach.py index bde3e62ea8..1f7e500969 100644 --- a/predicators/approaches/agent_program_world_model_approach.py +++ b/predicators/approaches/agent_program_world_model_approach.py @@ -27,12 +27,7 @@ import numpy as np from predicators import utils -from predicators.agent_sdk import learn_prompts -from predicators.agent_sdk.session_base import AgentSessionFatalError, \ - query_fatal_error from predicators.agent_sdk.tools import _SnapshotTarget -from predicators.agent_sdk.tools.digests import render_options_digest, \ - render_types_digest from predicators.agent_sdk.tools.program_synthesis import CandidateLoader, \ create_program_synthesis_tools from predicators.agent_sdk.tools.snapshots import finalize_versioned_snapshot @@ -41,9 +36,9 @@ AgentSimPredicateInventionApproach from predicators.code_sim_learning.program_world_model import \ ProgramOptionModel, ProgramWorldModel, load_program_world_model, \ - option_transitions, roll_program_latents + roll_program_latents from predicators.settings import CFG -from predicators.structs import LowLevelTrajectory, State, Task +from predicators.structs import LowLevelTrajectory, State logger = logging.getLogger(__name__) @@ -95,27 +90,6 @@ def _world_model_paths(paths: _SynthesisPaths) -> Dict[str, str]: # ── Learning ───────────────────────────────────────────────── - def _learn_simulator(self, trajectories: List[LowLevelTrajectory]) -> None: - """Run one program-synthesis session and deploy what it wrote.""" - self._fit_trajectories = list(trajectories) - self._persist_fit_trajectories("recorded") - usable = [ - t for t in trajectories if t.actions and t.actions[0].has_option() - ] - if not usable and not CFG.agent_sim_learn_zero_shot: - logger.warning("No skill-level transitions; skipping world " - "model synthesis.") - return - if not usable: - logger.info("Zero-shot synthesis: no skill-level transitions; " - "the agent writes the world model without data.") - program = self._run_program_synthesis_session(trajectories) - if program is None: - logger.warning("Synthesis produced no loadable world model; " - "the previous model stands.") - return - self._install_program(program) - def _install_program(self, program: ProgramWorldModel) -> None: self._program = program self._program_model = ProgramOptionModel(program, seed=CFG.seed) @@ -123,46 +97,6 @@ def _install_program(self, program: ProgramWorldModel) -> None: logger.info("Deployed the program world model (latent over %s).", dict(program.latent_features) or "nothing") - def _run_program_synthesis_session( - self, trajectories: List[LowLevelTrajectory] - ) -> Optional[ProgramWorldModel]: - paths = self._resolve_synthesis_paths() - wm_paths = self._world_model_paths(paths) - extra_paths = self._compute_extra_synthesis_paths(paths.base) - exec_ns = self._build_synthesis_exec_ns(trajectories) - self._attach_program_session_state(exec_ns, trajectories, paths, - wm_paths, extra_paths) - # Fresh session so the synthesis prompt + tools take effect. - self._close_agent_session() - self._ensure_agent_session() - structs_ref = self._write_structs_reference() - message = self._build_program_learn_message(trajectories, paths, - wm_paths, structs_ref, - extra_paths) - try: - responses = self._query_agent_sync(message, kind="learn") - dead = query_fatal_error(responses) - if dead is not None: - raise AgentSessionFatalError( - "The learn session died without the agent doing any " - f"work ({dead}); refusing to checkpoint this cycle as " - "learned.") - finally: - ctx = self._tool_context - ctx.extra_session_hooks = {} - ctx.extra_mcp_tools = [] - ctx.probe_artifact_loaders.clear() - ctx.probe_option_model_provider = None - ctx.probe_fit_provider = None - ctx.probe_validation_provider = None - ctx.probe_residuals_provider = None - ctx.probe_score_provider = None - ctx.probe_param_status = None - ctx.learn_cycle_index = None - self._learning_mode = False - self._close_agent_session() - return self._load_program_artifacts(wm_paths, extra_paths) - def _attach_program_session_state( self, exec_ns: Dict[str, Any], @@ -283,59 +217,6 @@ def _provider() -> ProgramOptionModel: return _provider - def _build_program_learn_message( - self, - trajectories: List[LowLevelTrajectory], - paths: _SynthesisPaths, - wm_paths: Dict[str, str], - structs_ref: str, - extra_paths: Dict[str, str], - ) -> str: - predicates = self._get_all_predicates() - n_trajs = len(trajectories) - n_demos = sum(1 for t in trajectories if t.is_demo) - n_transitions = sum( - len(option_transitions(t, predicates)) for t in trajectories) - prior: List[str] = [] - if os.path.isfile(wm_paths["world_model_file"]): - prior.append("`./world_model.py`") - if os.path.isfile(os.path.join(paths.base, "predicates.py")): - prior.append("`./predicates.py`") - session_tool_names = (self._agent_session.tool_names - if self._agent_session is not None else []) - extra_messages: List[str] = [] - if not trajectories and CFG.agent_sim_learn_zero_shot: - extra_messages.append( - learn_prompts.render_program_zero_shot_message()) - extra_message = self._extra_synthesis_message(extra_paths) - if extra_message: - extra_messages.append(extra_message) - return learn_prompts.build_program_learn_message( - n_trajs=n_trajs, - n_transitions=n_transitions, - n_demos=n_demos, - n_interaction=n_trajs - n_demos, - trajectory_listing=self._format_trajectory_listing(trajectories), - structs_ref=structs_ref, - predicate_listing=self._format_predicate_signatures(predicates), - types_digest=render_types_digest(self._tool_context.types), - options_digest=render_options_digest( - self._tool_context.options, - gt_options_ref_path=self._tool_context.gt_options_ref_path), - world_model_file=wm_paths["world_model_file_for_agent"], - objective_block=self._format_objective_block(), - prior_state_block=learn_prompts.render_prior_state_block(prior), - tools_block=learn_prompts.render_tools_block(session_tool_names), - extra_messages=extra_messages, - ) - - def _build_synthesis_system_prompt(self) -> str: - return learn_prompts.build_program_learn_system_prompt( - scene_viz_hint=self._scene_viz_hint(), - extra_sections=self._extra_synthesis_system_prompt_sections(), - workflow_extra=self._synthesis_workflow_extra(), - ) - def _load_program_artifacts( self, wm_paths: Dict[str, str], extra_paths: Dict[str, str]) -> Optional[ProgramWorldModel]: @@ -371,19 +252,6 @@ def _load_program_file( # ── Belief over the hidden state ───────────────────────────── - def _attach_initial_latent(self, task: Task) -> Task: - """Seed ``task.init.latent`` with the nominal particle (a seeded draw - from the program's ``initial_latent``).""" - if self._program_model is None: - return task - init_state = task.init.copy() - init_state.latent = self._program_model.initial_latent( - task.init, rng=np.random.default_rng(CFG.seed)) - return Task(init=init_state, - goal=task.goal, - alt_goal=task.alt_goal, - goal_nl=task.goal_nl) - def _belief_particles(self) -> List[Dict[str, float]]: """Distinct draws from ``initial_latent`` for the current task: the capture gate's margin points (empty until a model exists).""" diff --git a/predicators/approaches/agent_session_mixin.py b/predicators/approaches/agent_session_mixin.py index 14e9e4ad1c..c83dc77fee 100644 --- a/predicators/approaches/agent_session_mixin.py +++ b/predicators/approaches/agent_session_mixin.py @@ -1,8 +1,8 @@ """Mixin providing shared agent session infrastructure. Extracts common code for ToolContext initialization, lazy -AgentSessionManager creation, async-to-sync bridging, and agent explorer -creation shared by AgentModelFreeApproach and its subclasses. +AgentSessionManager creation and async-to-sync bridging shared by +AgentModelFreeApproach and its subclasses. """ import logging import os @@ -15,8 +15,6 @@ SessionManagerProtocol, run_async_sync, run_query_sync from predicators.agent_sdk.tools import ALL_TOOL_NAMES, ToolContext, \ create_mcp_tools, get_allowed_tool_list -from predicators.explorers import create_explorer -from predicators.explorers.base_explorer import BaseExplorer from predicators.settings import CFG from predicators.structs import ParameterizedOption, Predicate, Task, Type @@ -31,7 +29,7 @@ class AgentSessionMixin: And may optionally override: - _get_solve_tool_names() -- complete tool surface for - solve / explore sessions. May mix static MCP tool names with + play sessions. May mix static MCP tool names with names of dynamic ``SdkMcpTool`` instances. ``None`` = all static MCP tools, ``[]`` = none. - _get_synthesis_tool_names() -- complete tool surface for @@ -54,11 +52,6 @@ class AgentSessionMixin: # by the sim-learning approach around synthesis sessions; the class # default keeps plain solve-only hosts working without declaring it. _learning_mode: bool = False - # Flipped around explorer creation (``get_interaction_requests``) so - # explore sessions carry their own phase tag: their system prompt is - # saved as ``system_prompt_explore.md`` next to the solve and - # synthesis ones instead of overwriting the solve copy. - _explore_phase: bool = False # Phase tag the live ``_agent_session`` was created with; a query # under a different phase closes and rebuilds the session so the # saved prompt, tools, and CLAUDE.md always match the active phase. @@ -101,7 +94,7 @@ def _get_agent_system_prompt(self) -> str: raise NotImplementedError def _get_solve_tool_names(self) -> Optional[List[str]]: - """Return the complete tool surface for solve / explore sessions. + """Return the complete tool surface for play sessions. May mix static MCP tool names with names of dynamic ``SdkMcpTool`` instances. ``None`` means "all static MCP tools"; @@ -123,7 +116,7 @@ def _get_synthesis_tool_names(self) -> Optional[List[str]]: return [] def _get_sandbox_reference_files(self) -> Dict[str, str]: - """Return extra reference files for the docker sandbox. + """Return extra reference files for the sandbox. Maps destination paths (relative to ``/sandbox/reference/``) to source paths (relative to the repo root). Override in @@ -138,17 +131,12 @@ def _get_sandbox_reference_files(self) -> Dict[str, str]: def _ensure_agent_session(self) -> None: """Create the agent session manager if needed. - When ``SessionConfig.use_docker_sandbox`` is ``True``, creates a - ``DockerSessionManager`` that runs ``ClaudeSDKClient`` inside a - Docker container with full built-in tools (Bash, Read, Write, - …). Otherwise creates the normal in-process - ``AgentSessionManager``. + When ``SessionConfig.use_local_sandbox`` is ``True``, creates a + ``LocalSandboxSessionManager`` that runs the agent in a sandbox + directory with its built-in tools (Bash, Read, Write, ...). + Otherwise creates the in-process ``AgentSessionManager``. """ - phase = ("synthesis" if self._learning_mode else - "explore" if self._explore_phase else "solve") - # The tools read the phase for phase-dependent facts such as - # the real episode's step budget (explore episodes are capped by - # max_num_steps_interaction_request, tests by the horizon). + phase = "synthesis" if self._learning_mode else "solve" self._tool_context.phase = phase if self._agent_session is not None: if self._agent_session_phase == phase: @@ -220,21 +208,7 @@ def _ensure_agent_session(self) -> None: logger.info("\n".join(lines)) session: SessionManagerProtocol - if config.use_docker_sandbox: - from predicators.agent_sdk.docker_sandbox import \ - DockerSessionManager # pylint: disable=import-outside-toplevel - session = DockerSessionManager( - system_prompt=self._get_agent_system_prompt(), - log_dir=self._get_log_dir(), - model_name=config.model_name, - tool_context=self._tool_context, - tool_names=tool_names, - image=config.docker_image, - extra_reference_files=self._get_sandbox_reference_files(), - phase=phase, - config=config, - ) - elif config.use_local_sandbox: + if config.use_local_sandbox: from predicators.agent_sdk.local_sandbox import \ LocalSandboxSessionManager # pylint: disable=import-outside-toplevel session = LocalSandboxSessionManager( @@ -279,7 +253,7 @@ def _ensure_agent_session(self) -> None: # suffix) as full_system_prompt_{phase}.md during sandbox setup; for # in-process sessions the approach prompt IS the full prompt, so # save it under the same name here and nothing else. - if not (config.use_docker_sandbox or config.use_local_sandbox): + if not config.use_local_sandbox: log_dir = self._get_log_dir() os.makedirs(log_dir, exist_ok=True) prompt_path = os.path.join(log_dir, @@ -323,22 +297,3 @@ def _query_agent_sync(self, message: str, self._ensure_agent_session() assert self._agent_session is not None return run_query_sync(self._agent_session, message, **query_kwargs) - - def _create_agent_explorer( - self, - predicates: Set[Predicate], - options: Set[ParameterizedOption], - name: str = "agent_model_free", - ) -> BaseExplorer: - """Create an agent explorer with tool_context and agent_session.""" - self._ensure_agent_session() - return create_explorer( - name, - predicates, - options, - self._types, - self._action_space, - self._train_tasks, - tool_context=self._tool_context, - agent_session=self._agent_session, - ) diff --git a/predicators/approaches/agent_sim_learning_approach.py b/predicators/approaches/agent_sim_learning_approach.py index b1f3925825..1f910b482c 100644 --- a/predicators/approaches/agent_sim_learning_approach.py +++ b/predicators/approaches/agent_sim_learning_approach.py @@ -1,26 +1,17 @@ -"""Agent sim-learning approach: learns a simulator program online. - -Extends AgentModelBasedApproach to learn residual dynamics via an -agent-synthesized step-level simulator with parameterized process -rules. Parameters are fitted by Levenberg-Marquardt (fitting.py). - -The approach creates a base oracle (PyBullet with process -dynamics disabled) and composes it with the learned step-level -dynamics into a single simulator function, plugged into a standard -_OracleOptionModel for true per-step interleaving. - -Example command:: - - python predicators/main.py --env pybullet_boil \ - --approach agent_sim_learning --seed 0 \ - --num_train_tasks 10 --num_test_tasks 5 \ - --num_online_learning_cycles 5 --explorer agent_model_free +"""The model half of the EMPIRIC arms: the agent's simulator program, its +fitted parameters and the belief over them. + +The agent writes a simulator (a subclass of the base env, or residual +rules over it); the harness loads it, fits or deploys its parameters, +keeps their uncertainty, and composes it with a base oracle (PyBullet +with process dynamics disabled) into the option model plans are +rehearsed in. The continual arms (``agent_continual_approach``) drive +all of it from their play rounds. """ import copy import dataclasses import hashlib -import inspect import logging import math import os @@ -29,7 +20,6 @@ from typing import Any, Callable, Collection, ContextManager, Dict, \ FrozenSet, Iterator, List, Optional, Sequence, Set, Tuple -import dill as pkl import numpy as np import pybullet from gym.spaces import Box @@ -37,16 +27,13 @@ from predicators import utils from predicators.agent_sdk import learn_prompts from predicators.agent_sdk.fit_status import format_fit_status -from predicators.agent_sdk.session_base import AgentSessionFatalError, \ - max_session_log_number, query_fatal_error +from predicators.agent_sdk.session_base import max_session_log_number from predicators.agent_sdk.tools import SYNTHESIS_TOOL_NAMES, \ - _SnapshotTarget, create_synthesis_tools, evaluate_states_with, \ - finalize_versioned_snapshot, make_write_snapshot_hook -from predicators.agent_sdk.tools.digests import render_options_digest, \ - render_trajectory_digest, render_types_digest -from predicators.approaches.agent_model_based_approach import \ - AgentModelBasedApproach -from predicators.approaches.sampler_learning_mixin import SamplerLearningMixin + _SnapshotTarget, evaluate_states_with, finalize_versioned_snapshot, \ + make_write_snapshot_hook +from predicators.agent_sdk.tools.digests import render_trajectory_digest +from predicators.approaches.agent_model_free_approach import \ + AgentModelFreeApproach from predicators.approaches.synthesis_validation import \ build_candidate_option_model, carry_over_params from predicators.code_sim_learning.active_experiment import laplace_ensemble, \ @@ -57,18 +44,14 @@ from predicators.code_sim_learning.fit_space import FitResult, ParamSpec, \ declared_interval_fit_result, declared_interval_report from predicators.code_sim_learning.fitting import FIT_NOISE_SIGMA, \ - compute_sse, compute_sse_recurrent, fit_rule_parameters, \ - fit_rule_parameters_latent, log_param_changes, log_sse_breakdown -from predicators.code_sim_learning.identifiability import Verdict, \ - format_identifiability, physics_sigma_points + compute_sse, compute_sse_recurrent, log_sse_breakdown +from predicators.code_sim_learning.identifiability import physics_sigma_points from predicators.code_sim_learning.latent_tracker import LatentTracker, \ make_latent_tracker, make_subclass_latent_tracker from predicators.code_sim_learning.model_state import has_model_state -from predicators.code_sim_learning.orchestrator import \ - prior_parameter_belief, run_rollout_sysid +from predicators.code_sim_learning.orchestrator import prior_parameter_belief from predicators.code_sim_learning.parameter_belief import BeliefConfig, \ ParameterBelief, stable_seed -from predicators.code_sim_learning.physical_sysid import fit_params_rollout from predicators.code_sim_learning.rollout_env import RolloutTrajectory, \ dispose_env, physical_param_anchors from predicators.code_sim_learning.rollout_objective import compute_rollout_sse @@ -80,13 +63,12 @@ observation_view, read_latent_init, read_physical_param_specs, \ read_residual_env, read_simulator_components, stamp_physical_spec_scales from predicators.envs import create_new_env -from predicators.ground_truth_models import get_gt_simulator from predicators.observation_noise import ObservationNoise from predicators.option_model import _OptionModelBase, _OracleOptionModel from predicators.settings import CFG -from predicators.structs import Action, Dataset, DerivedPredicate, \ - GroundAtom, InteractionResult, LowLevelTrajectory, ParameterizedOption, \ - Predicate, State, Task, Type, step_option_labels +from predicators.structs import Action, DerivedPredicate, GroundAtom, \ + LowLevelTrajectory, ParameterizedOption, Predicate, State, Task, Type, \ + step_option_labels logger = logging.getLogger(__name__) @@ -193,23 +175,13 @@ def residual_hint_from_hits(hits: Dict[Tuple[str, str], int], return {t: sorted(fs) for t, fs in out.items()} -class AgentSimLearningApproach(SamplerLearningMixin, AgentModelBasedApproach): - """Bilevel planning with a learned step-level simulator. - - During online learning: - 1. Collect trajectories (inherited from AgentModelBasedApproach) - 2. Segment into option-level transitions - 3. Synthesize parameterized residual rules via Claude agent - 4. Fit rule parameters via Levenberg-Marquardt - 5. Compose with base oracle into a combined simulator - 6. Build _OracleOptionModel with the combined simulator - - During solving: - - Uses the learned model for plan validation in backtracking - refinement. +class AgentSimLearningApproach(AgentModelFreeApproach): + """The agent-written simulator, its parameters and their belief. - Per-skill sampler learning (mode resolution, synthesis session - plumbing, loading) lives in :class:`SamplerLearningMixin`. + Loads the simulator the agent writes, deploys its parameters (the + agent's published ``sim.fit`` or its declared values), keeps the + belief over them, and composes it with the base oracle into the + ``_OracleOptionModel`` the probe rehearses plans in. """ # Allowlist of env predicate names surfaced to the agent; None keeps @@ -342,7 +314,6 @@ def __init__(self, # (consumed in the next learn-phase prompt). self._current_simulator_version: Optional[str] = None self._current_predicates_version: Optional[str] = None - self._init_sampler_learning_state() # Partial-observability latent block: loaded from a simulator's # LATENT_INIT export (None ⇒ no latent state). When the loaded # rules use the recurrent 5-arg signature, fitting, the combined @@ -391,33 +362,6 @@ def __init__(self, # declaration/data signature); values are # orchestrator._FitComputation bundles. self._sysid_fit_cache: Dict[Tuple, Any] = {} - # Final per-cycle fit history for the cross-cycle consistency - # check: name -> (map_value, posterior_std_fit_space, scale). - # Mutually-incompatible confident fits across cycles are the - # signature of an overconfident probe; flagged, and the verdict - # downgraded, rather than silently trusted. - self._sysid_fit_history: Dict[str, Tuple[float, float, str]] = {} - # A rejected (INCONSISTENT) fit awaiting confirmation: - # name -> (map_value, posterior_std_fit_space). If the NEXT - # cycle's independent fit lands within the consistency band of - # the pending value, the jump is accepted as real (two - # independent fits agree); until then the trusted history value - # holds. Without this, a genuinely-updated fit would read - # INCONSISTENT against stale history forever. - self._sysid_pending_fit: Dict[str, Tuple[float, float]] = {} - # The applied physical params as of the last CYCLE-LEVEL fit - - # the reference the INCONSISTENT hold policy reverts to. - # Deliberately not _identified_physical_params: the agent's - # in-session sim.fit calls mutate that dict, so "hold the - # currently-applied value" was a no-op that held the very fit - # it refused to trust (run_20260724_232411 seed2 cycle 2: - # "holding the currently-applied 0.6267" - 0.6267 WAS the - # distrusted new fit, applied minutes earlier in-session). - self._cycle_applied_physical: Dict[str, float] = {} - # Agent-facing digest of the latest rollout fit (unexplainable - # segments, unidentified/insensitive params, cross-cycle - # conflicts); surfaced to the explorer as experiment objectives. - self._last_sysid_diagnostics: str = "" @classmethod def get_name(cls) -> str: @@ -459,27 +403,6 @@ def _compute_kept_initial_predicates(self) -> Set[Predicate]: # ── Agent session hooks ────────────────────────────────────── - def _get_agent_system_prompt(self) -> str: - if self._learning_mode: - return self._build_synthesis_system_prompt() - prompt = super()._get_agent_system_prompt() - base_sim_refs = self._base_sim_reference_paths() - if base_sim_refs: - ref_listing = "\n".join(f" - {r}" for r in base_sim_refs) - prompt += ( - "\n\n## Base Simulator Source\n" - "The environment simulator's own source code is " - "available (read-only):\n" - f"{ref_listing}\n" - "It covers the observable sim core: scene geometry and " - "constants, body construction, physics stepping, and " - "state read/write. It deliberately omits the hidden " - "domain-specific dynamics, task generation, and goal " - "semantics. Read it to ground your spatial and physical " - "reasoning (dimensions, contact geometry, actuation) " - "instead of guessing from images or trial and error.\n") - return prompt - def _get_sandbox_reference_files(self) -> Dict[str, str]: files = super()._get_sandbox_reference_files() # Base-sim source rides the standard reference channel so every @@ -519,20 +442,12 @@ def _get_synthesis_tool_names(self) -> Optional[List[str]]: # ── Subclass hooks ────────────────────────────────────────── # Default implementations are no-ops so subclasses can add - # predicate-invention (or other) extensions without copying - # _synthesize_with_agent. + # predicate-invention (or other) extensions. def _learning_cycle_index(self) -> int: - """0-based cycle index used in versioned snapshot filenames. - - Matches main.py's "ONLINE LEARNING CYCLE i" numbering exactly: - ``_online_learning_cycle`` is incremented before this class's - online simulator learn runs, so subtracting 1 recovers the - cycle the session belongs to. The offline (pre-cycle-0) learn - yields -1, which the snapshot/journal formatters render as - "offline" - keeping it distinct from cycle 0's online pass. - """ - return self._online_learning_cycle - 1 + """Index used in versioned snapshot filenames; the continual arms + number snapshots by level.""" + return 0 def _compute_extra_synthesis_paths(self, base: str) -> Dict[str, str]: """Return extra path bindings for the synthesis sandbox.""" @@ -555,34 +470,6 @@ def _install_extra_synthesis_surfaces( """ del exec_ns, base_pred_triples, inferred_hint, extra_paths - def _extra_synthesis_message(self, extra_paths: Dict[str, str]) -> str: - """Return text to append to the agent's first synthesis message. - - Under ``CFG.partially_observable`` this is the short partial- - observability note; subclasses that override MUST chain via - ``super()`` so the note survives. - """ - del extra_paths - if CFG.partially_observable: - return learn_prompts.render_partial_observability_message() - return "" - - def _extra_synthesis_system_prompt_sections(self) -> List[str]: - """Sections a subclass adds to the synthesis system prompt. - - Inserted after the validation guidance and before the recurrent - rules tutorial (partial observability) and the plan format. - Subclasses that override MUST chain via ``super()``. - """ - return [] - - def _extra_synthesis_latent_sections(self) -> List[str]: - """Sections a subclass adds after the recurrent-rules tutorial. - - Only rendered under ``CFG.partially_observable``. - """ - return [] - def _post_synthesis_loading( self, extra_paths: Dict[str, str], @@ -649,33 +536,6 @@ def _build_synthesis_session_hooks( # ── Learning ──────────────────────────────────────────────── - def learn_from_offline_dataset(self, dataset: Dataset) -> None: - super().learn_from_offline_dataset(dataset) - self._learn_simulator(self._get_all_trajectories()) - # The single post-offline checkpoint, AFTER the simulator learn - # (the base hook is a no-op for this class, see below). - self.save(None) - - def learn_from_interaction_results( - self, results: Sequence[InteractionResult]) -> None: - # Capture the index BEFORE super() increments it: the checkpoint - # below must be the one this cycle's filename denotes. - cycle = self._online_learning_cycle - super().learn_from_interaction_results(results) - self._learn_simulator(self._get_all_trajectories()) - # The single per-cycle checkpoint, AFTER this cycle's simulator - # learning, so a resume never re-pays a completed learn and never - # mistakes a pre-learn file for a completed cycle. (The base hook - # that would have saved pre-learn is a no-op for this class.) - self.save(cycle) - - def _checkpoint_after_offline_learning(self) -> None: - """No-op: this class checkpoints after its own simulator learn.""" - - def _checkpoint_after_interaction_results(self, cycle: int) -> None: - """No-op: this class checkpoints after its own simulator learn.""" - del cycle - # ── Checkpointing ──────────────────────────────────────────── # The base checkpoint (AgentModelFreeApproach.save/load) persists # the datasets + cycle counter. This approach's real state is split @@ -684,21 +544,19 @@ def _checkpoint_after_interaction_results(self, cycle: int) -> None: # which are embedded as file CONTENTS - run dirs are minted per run # and pruned, so a path reference to the old run's sandbox would be # fragile. Closures (_residual_rules, _learned_simulator, the option - # model, learned predicates/samplers) are never pickled: they are + # model, learned predicates) are never pickled: they are # rebuilt from the restored files in _rehydrate_from_artifacts. _save_suffix: str = "AgentSimLearner" _CHECKPOINT_SANDBOX_FILES: Tuple[str, ...] = ("simulator.py", "predicates.py", - "samplers.py", "ground_samplers.py", "notes.md", "journal.md", "attempts.md", "strategy.md", "open_questions.md") _CHECKPOINT_SANDBOX_DIRS: Tuple[str, ...] = ("simulator_versions", - "predicates_versions", - "samplers_versions") + "predicates_versions") _CHECKPOINT_MAX_FILE_BYTES = 2 * 1024 * 1024 def _checkpoint_sandbox_dir(self) -> str: @@ -783,16 +641,12 @@ def _extra_save_state(self) -> Dict[str, Any]: dict(self._fit_evidence_history), "identified_physical_sigma_points": list(self._identified_physical_sigma_points), - "sysid_fit_history": - dict(self._sysid_fit_history), "residual_features": dict(self._residual_features), "current_simulator_version": self._current_simulator_version, "current_predicates_version": self._current_predicates_version, - "current_samplers_version": - self._current_samplers_version, "sandbox_files": self._collect_sandbox_artifacts(), "git_describe": @@ -812,8 +666,8 @@ def _load_extra_save_state(self, save_dict: Dict[str, Any]) -> None: "is at %s - resuming across code versions is untested.", saved_rev, current_rev) self._resume_query_count = int(save_dict.get("agent_query_count", 0)) - # In-place update: _ParamsView holders (invented predicate and - # sampler closures) alias this exact dict object. + # In-place update: _ParamsView holders (invented predicate + # closures) alias this exact dict object. self._fitted_params.clear() self._fitted_params.update(save_dict.get("fitted_params") or {}) self._fit_sse = save_dict.get("fit_sse", float("inf")) @@ -831,18 +685,12 @@ def _load_extra_save_state(self, save_dict: Dict[str, Any]) -> None: save_dict.get("carried_physical_prior") or {}) self._fit_evidence_history = dict( save_dict.get("fit_evidence_history") or {}) - self._sysid_fit_history = dict( - save_dict.get("sysid_fit_history") or {}) self._residual_features = dict( save_dict.get("residual_features") or {}) self._current_simulator_version = save_dict.get( "current_simulator_version") self._current_predicates_version = save_dict.get( "current_predicates_version") - # pylint: disable-next=attribute-defined-outside-init - # (initialized by SamplerLearningMixin's init hook) - self._current_samplers_version = save_dict.get( - "current_samplers_version") self._restore_sandbox_artifacts(save_dict.get("sandbox_files") or {}) self._rehydrate_from_artifacts() # AFTER rehydration: _apply_identified_physical_params clears @@ -860,8 +708,7 @@ def _rehydrate_from_artifacts(self) -> None: Order matters: simulator.py first (rules + latent init + physical specs), then the option model, then identified physics onto the base env, then subclass artifacts (predicates read the - already- restored ``_fitted_params``), then samplers and the - ensemble. + already- restored ``_fitted_params``), then the ensemble. """ paths = self._resolve_synthesis_paths() if not os.path.isfile(paths.simulator_file): @@ -934,126 +781,12 @@ def _step_fn(s: State, c: Any) -> Any: self._apply_identified_physical_params( self._identified_physical_params) self._rehydrate_extra_artifacts(paths.base) - if self._samplers_enabled(): - sampler_paths = self._sampler_paths(paths.base) - self._synthesized_samplers = self._load_samplers_from_module_file( - sampler_paths["samplers_file"]) self._rebuild_param_ensemble() logger.info( "Rehydrated learned simulator from checkpoint artifacts " - "(%d rules, %d fitted params, %d learned predicates, " - "%d samplers).", len(rules), len(self._fitted_params), - len(getattr(self, "_learned_predicates", set()) or set()), - len(self._synthesized_samplers)) - - def _learn_simulator(self, trajectories: List[LowLevelTrajectory]) -> None: - """Synthesize rules, fit parameters, and build the option model.""" - # Cache for recurrent fitting: lets _group_triples_by_trajectory - # slice the flat base_pred_triples back into per-trajectory chunks - # (latent threads within a trajectory, not across). Harmless for - # fully-observable (legacy) simulators, which never regroup. - self._fit_trajectories = list(trajectories) - # Dumped HERE, where the data arrives, rather than only inside the - # sysID fit: a cycle where the agent declines to fit is exactly the - # one worth post-morteming, and that is the branch that never ran. - # run_20260817_171402 declined on a sweep that returned one identical - # SSE for every value of five parameters, and left nothing on disk to - # explain it -- the episode had to be written off. - self._persist_fit_trajectories("recorded") - # New data invalidates the memoized explainability verdicts and - # the memoized whole fits. - self._explainability_cache.clear() - self._sysid_fit_cache.clear() - # Decide how samplers are obtained this cycle: ground-truth (if - # requested and available for the env) else agent synthesis. GT - # samplers are static, so install them up front, independent of - # whether simulator learning runs below (it is skipped when there - # are no step transitions and no oracle sim program to fall - # back on, e.g. when every demo failed). - self._maybe_install_oracle_samplers() - # Two parallel triple lists drive the rest of this method: - # * obs_triples - raw (s_t, a, s_{t+1}) from the data. - # * base_pred_triples - same triples but s_t replaced by the - # base sim's one-step prediction. The rules run on top of that - # prediction; SSE compares against s_{t+1}. - obs_triples = self._extract_obs_triples(trajectories) - if (not obs_triples and not CFG.agent_sim_learn_oracle_sim_program - and not CFG.agent_sim_learn_zero_shot): - logger.warning("No step transitions; skipping simulator learning.") - return - if obs_triples: - # Headless env for the pre-compute: reusing the GUI base_env - # corrupts its visual-shape state after a few hundred steps. - fit_env = create_new_env(CFG.env, - do_cache=False, - use_gui=False, - skip_residual_dynamics=True) - logger.info("Pre-computing base states for %d transitions.", - len(obs_triples)) - try: - base_pred_triples = self._compute_base_pred_triples( - obs_triples, fit_env) - finally: - # This env is rebuilt every learning cycle; dispose it - # (main client AND any secondary probe world) or each - # cycle leaks a full physics world (~145MB for the - # domino env). - dispose_env(fit_env) - inferred_hint = self._infer_residual_features_from_scan( - obs_triples, base_pred_triples) - logger.info("Residual features (data-driven hint): %s", - inferred_hint) - elif CFG.agent_sim_learn_oracle_sim_program: - # The oracle sim program is data-free (rules and parameter - # inits come from get_gt_simulator), so a run whose every - # demo failed still gets a working option model; the fit - # below degrades to the declared inits. - logger.warning("No step transitions; loading oracle sim " - "program without data.") - base_pred_triples = [] - inferred_hint = {} - else: - # Zero-shot synthesis (ablation A2): the session runs with - # nothing recorded, so the artifacts come from the task - # description, the scene and the agent's own knowledge; the - # params deploy at their declared inits. - logger.info("Zero-shot synthesis: no step transitions; the " - "agent writes its artifacts without data.") - base_pred_triples = [] - inferred_hint = {} - - self._synthesize_with_agent(trajectories, obs_triples, - base_pred_triples, inferred_hint) - - if self._residual_rules is not None and self._fitted_params: - rules, params = self._residual_rules, self._fitted_params - self._learned_simulator = LearnedSimulator( - step_fn=lambda s, c, _r=rules, _p=params: # type: ignore[misc] - apply_rules(s, _r, _p, cmds=c), - name="agent_synthesized") - elif self._learned_simulator is None: - logger.warning("Synthesis produced no simulator, skipping.") - return - - combined_sim = self._build_combined_simulator(self._learned_simulator) - self._option_model = self._build_option_model(combined_sim) - logger.info("Built learned option model (SSE: %.6f).", self._fit_sse) - - # When the simulator came from the oracle short-circuit no agent - # session ran above, so per-skill samplers (if enabled) get their - # own session here, after the option model is built so the - # session's probe (sim.refine) has a working simulator. When - # the agent *did* synthesize the simulator, samplers already rode - # along in that session and this is skipped. - if self._do_synthesize_samplers and \ - CFG.agent_sim_learn_oracle_sim_program: - if base_pred_triples: - self._synthesize_samplers_standalone(trajectories, - base_pred_triples, - inferred_hint) - else: - logger.warning("No step transitions; skipping standalone " - "sampler synthesis.") + "(%d rules, %d fitted params, %d learned predicates).", len(rules), + len(self._fitted_params), + len(getattr(self, "_learned_predicates", set()) or set())) def _build_option_model( self, @@ -1305,8 +1038,6 @@ def _make_candidate_probe_model_provider( self, simulator_file: str, trajectories: List[LowLevelTrajectory], - base_pred_triples: List[Tuple[State, Action, State]], - inferred_hint: Dict[str, List[str]], ) -> Callable[[], _OracleOptionModel]: """Lazy option-model builder behind the synthesis run_python. @@ -1342,16 +1073,13 @@ def _provider() -> _OracleOptionModel: if cache.get("digest") == digest: return cache["model"] fit_state = self._probe_fit_state() - rules, specs, features, ns = \ - self._load_simulator_from_module_file( - simulator_file, trajectories) + rules, specs, _, ns = self._load_simulator_from_module_file( + simulator_file, trajectories) if rules is None or specs is None: raise RuntimeError( "run_python probe: ./simulator.py failed to load " "(exec error or missing simulator exports) - " "fix the file and probe again.") - residual_features = (features - if features is not None else inferred_hint) # The candidate's rule parameters, which the joint belief # covers beside AGENT_PARAM_SPECS (see parameter_belief). setattr(self, "_probe_rule_specs", list(specs)) @@ -1378,14 +1106,8 @@ def _provider() -> _OracleOptionModel: # runs a fit it may not afford (sketch seed1 learn 011: an # implicit refit inside a probe hit the call cap and came # back as an empty "param fitting failed:"). - model, params, _ = build_candidate_option_model( - self, - rules, - specs, - residual_features, - base_pred_triples, - latent_init=latent_init, - fit=False) + model, params = build_candidate_option_model( + self, rules, specs, latent_init=latent_init) if CFG.agent_sim_learn_declared_params_only: status = ("at the DECLARED values of the current " "simulator.py (harness parameter estimation is " @@ -1411,19 +1133,6 @@ def _provider() -> _OracleOptionModel: # ── Active-experiment ensemble (info-seeking exploration) ──── - def _info_seeking_active(self) -> bool: - """Whether the proactive info-seeking apparatus should run now. - - Delegates to the run context's adaptive gate. Partial unit-test - objects have no ``_tool_context``; there, fall back to the plain - flag (adaptive gating needs the run-scoped refusal signal the - context carries). - """ - ctx = getattr(self, "_tool_context", None) - if ctx is not None: - return ctx.info_seeking_active() - return CFG.agent_explorer_info_seeking - def _rebuild_param_ensemble(self) -> None: """Rebuild the learned model's rule-parameter ensemble. @@ -1459,16 +1168,6 @@ def _rebuild_param_ensemble(self) -> None: or not CFG.agent_sim_learn_param_uncertainty): self._param_ensemble = [] return - if CFG.agent_sim_learn_oracle_sim_params: - # Oracle params carry no uncertainty: no fit ran, so the only - # ensemble on offer would be box jitter around the truth, - # which manufactures wrong models (a zero rate, a rewired - # lamp) that the capture gate would then demand every plan - # survive. Nothing to hedge against, so no ensemble. - self._param_ensemble = [] - logger.info("Oracle sim params: no rule-parameter ensemble " - "(nothing uncertain to sweep).") - return num_members = CFG.agent_explorer_info_ensemble_size self._param_ensemble, method = self._select_param_ensemble(num_members) logger.info( @@ -1536,9 +1235,8 @@ def score_atom_disagreement(self, state: State, i.e. an informative experiment. Returns 0.0 when the ensemble is trivial (<=1 member) or no atoms are given. - Wired into refinement as the info-scorer for the agent_model_based - explorer; a read-only query that leaves ``_fitted_params`` - unchanged on return. + Wired into the probe's refinement as the info-scorer; a read-only + query that leaves ``_fitted_params`` unchanged on return. Under ``agent_explorer_info_seeking_noise_aware`` with a declared observation-noise channel, each member reads the atoms from @@ -1621,153 +1319,6 @@ def _rule_param_override_scope( # ── Agent-based synthesis ──────────────────────────────────── - def _synthesize_with_agent( - self, - trajectories: List[LowLevelTrajectory], - obs_triples: List[Tuple[State, Action, State]], - base_pred_triples: List[Tuple[State, Action, State]], - inferred_hint: Dict[str, List[str]], - ) -> None: - """Obtain RESIDUAL_RULES / PARAM_SPECS / RESIDUAL_FEATURES, then fit. - - ``inferred_hint`` is passed to the agent as a starting point and - used as the eval/test scope until it declares its own - ``RESIDUAL_FEATURES``. CFG flag - ``agent_sim_learn_oracle_sim_program`` short-circuits the agent - session by loading the GT simulator instead (and - ``agent_sim_learn_oracle_sim_params`` additionally skips the - parameter fit; see :meth:`_fit_params_after_synthesis`). - """ - if CFG.agent_sim_learn_oracle_sim_program: - rules, specs, residual_features = \ - self._load_oracle_sim_program(inferred_hint) - else: - loaded = self._run_agent_synthesis_session(trajectories, - obs_triples, - base_pred_triples, - inferred_hint) - if loaded is None: - return - rules, specs, residual_features = loaded - self._residual_rules = rules - self._residual_features = residual_features - self._fit_params_after_synthesis(rules, specs, base_pred_triples, - residual_features) - - def _load_oracle_sim_program( - self, inferred_hint: Dict[str, List[str]] - ) -> Tuple[List, List[ParamSpec], Dict[str, List[str]]]: - """Load the ground-truth simulator instead of running an agent. - - ``get_gt_simulator`` dispatches by observability: in - partially-observable mode it returns the PO GT simulator - (gt_simulator_po.py - latent heat threaded across steps, - surfaced as the observable bubbling_level), which predicts only - observable features; otherwise it returns the fully-observable - gt_simulator.py (which reads/writes heat_level as a State - feature). The two factories gate on CFG.partially_observable so - the env-name dispatch resolves to exactly one module per run. - - Unless ``agent_sim_learn_oracle_sim_params`` also holds, the - declared parameter inits are perturbed so the subsequent fit - starts from a miscalibrated - not oracle - belief. - """ - rules, specs, residual_features = get_gt_simulator(CFG.env) - self._log_feature_set_diff(inferred_hint, residual_features, - "inferred", "oracle") - if not CFG.agent_sim_learn_oracle_sim_params: - specs = self._perturb_spec_inits(specs) - logger.info("Loaded oracle sim program (%d rules, %d params).", - len(rules), len(specs)) - return rules, specs, residual_features - - @staticmethod - def _perturb_spec_inits(specs: List[ParamSpec]) -> List[ParamSpec]: - """Perturb each spec's init with multiplicative Gaussian noise. - - Used when the oracle sim PROGRAM is loaded but its param VALUES - must still be learned: the fit then starts from a plausible but - wrong belief instead of the answer. Each perturbed init is - clipped to its spec's box. - """ - rng = np.random.default_rng(CFG.seed) - noise_scale = CFG.agent_sim_learn_oracle_sim_param_noise_scale - if noise_scale < 0.0: - raise ValueError("agent_sim_learn_oracle_sim_param_noise_scale " - "must be non-negative.") - perturbed = [] - for s in specs: - val = float( - np.clip(s.init_value * (1.0 + rng.normal(0, noise_scale)), - s.lo, s.hi)) - perturbed.append( - ParamSpec(s.name, - val, - lo=s.lo, - hi=s.hi, - scale=getattr(s, "scale", "linear"))) - return perturbed - - def _run_agent_synthesis_session( - self, - trajectories: List[LowLevelTrajectory], - obs_triples: List[Tuple[State, Action, State]], - base_pred_triples: List[Tuple[State, Action, State]], - inferred_hint: Dict[str, List[str]], - ) -> Optional[Tuple[List, List[ParamSpec], Dict[str, List[str]]]]: - """Run one agent synthesis session and load what it committed. - - Returns ``(rules, specs, residual_features)``, or None when the - session left no loadable simulator artifact. Per-skill samplers - (when enabled) ride along in the same session. - """ - paths = self._resolve_synthesis_paths() - extra_paths = self._compute_extra_synthesis_paths(paths.base) - sampler_paths = (self._sampler_paths(paths.base) - if self._do_synthesize_samplers else {}) - exec_ns = self._build_synthesis_exec_ns(trajectories) - self._attach_synthesis_session_state(exec_ns, trajectories, - base_pred_triples, inferred_hint, - paths, extra_paths, sampler_paths) - # Fresh session so the synthesis prompt + tools take effect. - self._close_agent_session() - self._ensure_agent_session() - structs_ref = self._write_structs_reference() - base_sim_refs = self._base_sim_reference_paths() - message = self._build_synthesis_learn_message( - trajectories, obs_triples, inferred_hint, paths, structs_ref, - extra_paths, sampler_paths, base_sim_refs) - try: - responses = self._query_agent_sync(message, kind="learn") - dead = query_fatal_error(responses) - if dead is not None: - # The synthesis session never ran (usage limit, auth, - # transport): nothing was learned, so this cycle must - # not be checkpointed as learned. The cycle's explore - # episodes are stashed (main._save_inflight_interactions), - # so a relaunch resumes at exactly this learn. Silently - # continuing once wrote a byte-identical checkpoint and - # burned a whole cycle (2026-08-27 run_20260827_121111). - raise AgentSessionFatalError( - "The learn session died without the agent doing any " - f"work ({dead}); refusing to checkpoint this cycle as " - "learned.") - finally: - self._tool_context.extra_session_hooks = {} - self._tool_context.extra_mcp_tools = [] - self._tool_context.probe_artifact_loaders.clear() - self._tool_context.probe_option_model_provider = None - self._tool_context.probe_fit_provider = None - self._tool_context.probe_validation_provider = None - self._tool_context.probe_param_status = None - self._tool_context.probe_residuals_provider = None - self._tool_context.learn_cycle_index = None - self._learning_mode = False - self._close_agent_session() - return self._load_synthesis_artifacts(trajectories, inferred_hint, - paths, extra_paths, - sampler_paths) - def _resolve_synthesis_paths(self) -> _SynthesisPaths: """Host- and agent-visible paths for one synthesis session. @@ -1776,10 +1327,9 @@ def _resolve_synthesis_paths(self) -> _SynthesisPaths: in __init__, but it isn't constructed until ``_ensure_agent_session()`` runs later in the session setup. - The agent-visible paths differ by sandbox backend: cwd-relative - for local-sandbox (the validation hook resolves against cwd and - rejects literal ``/sandbox/...`` paths), the docker mount point - for docker, the absolute host path otherwise. + The agent-visible paths are cwd-relative in the local sandbox + (the validation hook resolves against cwd and rejects literal + ``/sandbox/...`` paths) and absolute host paths otherwise. """ if CFG.agent_sdk_use_local_sandbox: sandbox_dir: Optional[str] = os.path.abspath( @@ -1791,9 +1341,6 @@ def _resolve_synthesis_paths(self) -> _SynthesisPaths: if CFG.agent_sdk_use_local_sandbox: simulator_file_for_agent = "./simulator.py" sandbox_dir_for_agent: Optional[str] = "." - elif sandbox_dir: - simulator_file_for_agent = "/sandbox/simulator.py" - sandbox_dir_for_agent = "/sandbox" else: simulator_file_for_agent = simulator_file sandbox_dir_for_agent = None @@ -1847,174 +1394,12 @@ def describe_trajectory(traj_idx: int, self._make_evaluate_trajectory_fn() return exec_ns - def _attach_synthesis_session_state( - self, - exec_ns: Dict[str, Any], - trajectories: List[LowLevelTrajectory], - base_pred_triples: List[Tuple[State, Action, State]], - inferred_hint: Dict[str, List[str]], - paths: _SynthesisPaths, - extra_paths: Dict[str, str], - sampler_paths: Dict[str, str], - ) -> None: - """Install this synthesis session's state on the tool context. - - Everything installed here is cleared by the caller's ``finally`` - once the session query returns. - """ - # Label tool output (e.g. attempt-log headers) with the - # learning cycle for the duration of this session. - self._tool_context.learn_cycle_index = self._learning_cycle_index() - # Build dynamic synthesis tools and attach them to the tool - # context *before* opening the session. The attached set is - # filtered against ``_get_synthesis_tool_names`` so that method - # is the single source of truth for what the agent sees: - # anything a builder constructs but the names list omits is - # dropped here. - # pylint: disable-next=import-outside-toplevel - from predicators.agent_sdk.belief_probe import _check_time_budget - toolkit = create_synthesis_tools( - exec_ns, - base_pred_triples, - inferred_hint, - simulator_file=paths.simulator_file, - versions_dir=paths.versions_dir, - approach=self, - sandbox_dir=paths.base, - sandbox_dir_for_agent=paths.sandbox_dir_for_agent, - cycle_index_provider=self._learning_cycle_index, - budget_check=lambda: _check_time_budget(self._tool_context), - ) - tools = list(toolkit.tools) - self._install_extra_synthesis_surfaces(exec_ns, base_pred_triples, - inferred_hint, extra_paths) - if self._do_synthesize_samplers: - self._install_sampler_surface(sampler_paths) - declared = set(self._get_synthesis_tool_names() or ()) - self._tool_context.extra_mcp_tools = [ - t for t in tools if getattr(t, "name", "") in declared - ] - # Point the probe at the CANDIDATE simulator for this session - # (never the stale pre-synthesis option model; on cycle 1 that - # wraps the real env), then merge the probe facade into - # run_python's namespace: synthesis sessions offer ONE exec - # namespace, so helpers defined next to the data are visible to - # probe sweeps (create_mcp_tools skips the solve-phase instance - # when this one is attached). Unconditional: with fit / refine / - # forward-validation all living on ``sim``, the probe IS the - # validation surface, so a synthesis session without it would - # have no way to test what it writes. Only ``sim``/``BeliefProbe`` - # are taken from the probe namespace: ``trajectories`` already - # binds the fit list and solve-only extras do not apply. - self._tool_context.probe_option_model_provider = \ - self._make_candidate_probe_model_provider( - paths.simulator_file, trajectories, base_pred_triples, - inferred_hint) - self._tool_context.probe_fit_provider = toolkit.fit_runner - self._tool_context.probe_validation_provider = toolkit.validation_runner - self._tool_context.probe_residuals_provider = \ - toolkit.residuals_runner - # pylint: disable-next=import-outside-toplevel - from predicators.agent_sdk.belief_probe import build_probe_namespace - probe_ns = build_probe_namespace(self._tool_context) - exec_ns["sim"] = probe_ns["sim"] - exec_ns["BeliefProbe"] = probe_ns["BeliefProbe"] - self._learning_mode = True - # PostToolUse hook: snapshot simulator.py / predicates.py on - # every successful Write/Edit/MultiEdit, so the version history - # covers everything the agent committed to file (not just - # states that happened to coincide with an eval call). Only - # active for this synthesis session. - snapshot_targets = self._build_write_snapshot_targets( - paths.simulator_file, paths.versions_dir, extra_paths) - if self._do_synthesize_samplers: - snapshot_targets.append( - self._sampler_snapshot_target(sampler_paths)) - self._tool_context.extra_session_hooks = ( - self._build_synthesis_session_hooks(snapshot_targets, paths.base)) - - def _build_synthesis_learn_message( - self, - trajectories: List[LowLevelTrajectory], - obs_triples: List[Tuple[State, Action, State]], - inferred_hint: Dict[str, List[str]], - paths: _SynthesisPaths, - structs_ref: str, - extra_paths: Dict[str, str], - sampler_paths: Dict[str, str], - base_sim_refs: Optional[List[str]] = None, - ) -> str: - """Compose the synthesis session's first user message. - - Gathers this cycle's data roster, digests, and reports and - renders ``learn_message.md``. Reads the just-opened session's - tool names, so the session must be open before this is called. - """ - n_trajs = len(trajectories) - n_demos = sum(1 for t in trajectories if t.is_demo) - # Start-of-session divergence report: when a prior model exists, - # score it (params refit to ALL data, so what remains is the - # structural gap) before the agent's first turn - the session - # then starts from "here is where the model breaks" instead of - # spending turns rediscovering it. The same report stays callable - # as `sim.residuals()` against every subsequent edit. With no - # prior model the "prior" is the bare base simulator and every - # mismatch is an unmodeled mechanism. - prior_state_block = self._format_prior_state_block(paths.base) - divergence_block = "" - if (self._tool_context.probe_residuals_provider is not None - and obs_triples): - try: - report = self._tool_context.probe_residuals_provider( - max_transitions=100000, - fit_params=not CFG.agent_sim_learn_declared_params_only) - divergence_block = learn_prompts.render_divergence_block( - report, has_prior_model=bool(prior_state_block)) - except Exception as e: # pylint: disable=broad-except - logger.warning("Skipping start-of-session residual report: %s", - e) - session_tool_names = (self._agent_session.tool_names - if self._agent_session is not None else []) - extra_messages = [] - if not trajectories and CFG.agent_sim_learn_zero_shot: - extra_messages.append(learn_prompts.render_zero_shot_message()) - extra_message = self._extra_synthesis_message(extra_paths) - if extra_message: - extra_messages.append(extra_message) - if self._do_synthesize_samplers: - extra_messages.append( - self._sampler_synthesis_message(sampler_paths)) - return learn_prompts.build_learn_message( - n_trajs=n_trajs, - n_transitions=len(obs_triples), - n_demos=n_demos, - n_interaction=n_trajs - n_demos, - trajectory_listing=self._format_trajectory_listing(trajectories), - structs_ref=structs_ref, - inferred_hint=str(inferred_hint), - predicate_listing=self._format_predicate_signatures( - self._get_all_predicates()), - types_digest=render_types_digest(self._tool_context.types), - options_digest=render_options_digest( - self._tool_context.options, - gt_options_ref_path=self._tool_context.gt_options_ref_path), - simulator_file=paths.simulator_file_for_agent, - objective_block=self._format_objective_block(), - prior_state_block=prior_state_block, - divergence_block=divergence_block, - base_sim_block=learn_prompts.render_base_sim_block(base_sim_refs - or []), - tools_block=learn_prompts.render_tools_block(session_tool_names), - extra_messages=extra_messages, - ) - def _load_synthesis_artifacts( self, trajectories: List[LowLevelTrajectory], inferred_hint: Dict[str, List[str]], paths: _SynthesisPaths, extra_paths: Dict[str, str], - sampler_paths: Dict[str, str], ) -> Optional[Tuple[List, List[ParamSpec], Dict[str, List[str]]]]: """Load the artifacts the finished session committed to disk. @@ -2078,8 +1463,6 @@ def _load_synthesis_artifacts( logger.info("Agent synthesized %d rules, %d params.", len(rules), len(specs)) self._post_synthesis_loading(extra_paths, specs) - if self._do_synthesize_samplers: - self._finalize_and_load_samplers(sampler_paths) return rules, specs, residual_features def _fit_params_after_synthesis( @@ -2089,30 +1472,12 @@ def _fit_params_after_synthesis( base_pred_triples: List[Tuple[State, Action, State]], residual_features: Dict[str, List[str]], ) -> None: - """Deploy the agent's parameters; only oracle programs fit here.""" + """Deploy the agent's parameters; the harness never fits here.""" if getattr(self, "_residual_env_cls", None) is not None and \ not specs and not self._physical_param_specs: self._fitted_params.clear() self._last_fit_result = None self._fit_sse = float("inf") - elif CFG.agent_sim_learn_oracle_sim_params: - self._fitted_params.clear() - self._fitted_params.update({s.name: s.init_value for s in specs}) - if self._physical_param_specs: - # Oracle mode: trust the agent-declared physical inits. - self._apply_identified_physical_params( - {s.name: s.init_value - for s in self._physical_param_specs}) - # No fit ran; the ensemble falls back to uniform perturbation. - self._last_fit_result = None - if base_pred_triples: - self._fit_sse = self._oracle_param_sse(rules, - base_pred_triples, - residual_features, - FIT_NOISE_SIGMA) - else: - logger.info("No transitions; skipping oracle-param SSE.") - self._fit_sse = float("inf") elif CFG.agent_sim_learn_declared_params_only: self._deploy_declared_params(rules, specs, base_pred_triples, residual_features) @@ -2158,33 +1523,14 @@ def _fit_params_after_synthesis( version, len(expected), self._fit_sse) applied = self._probe_fit_state().get("applied_physical") if self._physical_param_specs and applied: - # Mirror _fit_parameters_joint_rollout's deploy: the - # cycle-level applied snapshot and the physics-margin - # sigma points come from the published fit (applying - # resets the points, so set them after). + # The physics-margin sigma points come from the + # published fit (applying resets the points, so set + # them after). self._apply_identified_physical_params(dict(applied)) - self._cycle_applied_physical = dict(applied) if CFG.agent_sim_learn_param_uncertainty and int( CFG.belief_joint_draws) <= 0: self._identified_physical_sigma_points = list( self._probe_fit_state().get("sigma_points") or []) - elif CFG.agent_sim_learn_oracle_sim_program and base_pred_triples: - # This baseline supplies a program without an agent - # session. Fitting is its explicitly configured protocol. - logger.info("Oracle sim program: fitting its " - "parameters on the harness side.") - if self._physical_param_specs or has_physics_rules(rules): - fit_result, self._fit_sse = ( - self._fit_parameters_joint_rollout( - rules, specs, residual_features)) - elif has_latent_rules(rules): - fit_result, self._fit_sse = \ - self._fit_parameters_recurrent( - rules, specs, base_pred_triples, - residual_features) - else: - fit_result, self._fit_sse = fit_rule_parameters( - rules, specs, base_pred_triples, residual_features) else: self._deploy_unfitted_params(specs) if fit_result is not None: @@ -2212,7 +1558,6 @@ def _deploy_unfitted_params(self, specs: List[ParamSpec]) -> None: if physical or self._identified_physical_params: self._apply_identified_physical_params(physical) self._identified_physical_sigma_points = [] - self._cycle_applied_physical = dict(physical) # Retain the historical published fit for provenance and file # reversions, but do not use its SSE or posterior for this model. self._last_fit_result = None @@ -2268,7 +1613,6 @@ def _deploy_declared_params( if physical_specs: applied = {s.name: s.init_value for s in physical_specs} self._apply_identified_physical_params(applied) - self._cycle_applied_physical = dict(applied) self._identified_physical_sigma_points = \ self._physics_margin_points( applied, declared_interval_report(physical_specs), @@ -2456,47 +1800,6 @@ def _rollout_fit_trajectories( "keeping the whole trajectories.") return rollouts - def _persist_fit_trajectories(self, label: str = "fitted") -> None: - """Dump the raw fit trajectories for offline post-mortems. - - ``label`` distinguishes the two moments this is called from: - ``recorded`` when a cycle's data arrives (always), ``fitted`` - when a sysID fit has just run and the payload's identified - params mean something. The first is what makes a cycle that - declined to fit replayable at all. - - The rollout sysID fit data otherwise exists only in memory: - when run_20260724_232411 shipped friction fits 2-7 sigma from - the truth, the failing fits could not be replayed offline - the - episodes had to be approximately re-executed from logged plans, - which cannot reproduce mid-episode replans or the warm-env - recording context (exactly the suspected corruption channel). - One pickle per cycle-level fit under ``/fit_data/``; - never raises - persistence must not take down a run. - """ - if not CFG.code_sim_learning_persist_fit_data: - return - try: - out_dir = os.path.join(self._get_log_dir(), "fit_data") - os.makedirs(out_dir, exist_ok=True) - idx = len([f for f in os.listdir(out_dir) if f.endswith(".pkl")]) - path = os.path.join(out_dir, - f"fit_trajectories_{idx:03d}_{label}.pkl") - payload = { - "trajectories": - list(self._fit_trajectories), - "physical_param_specs": - list(self._physical_param_specs), - "identified_physical_params": - dict(self._identified_physical_params), - } - with open(path, "wb") as f: - pkl.dump(payload, f) - logger.info("Persisted %d fit trajectories to %s", - len(self._fit_trajectories), path) - except Exception as e: # pylint: disable=broad-except - logger.warning("Could not persist fit trajectories: %s", e) - def _apply_identified_physical_params( self, identified: Dict[str, float]) -> None: """Publish identified physical params into the planning base env. @@ -2591,392 +1894,8 @@ def previous_fit_evidence( } return None - def _fit_parameters_joint_rollout( - self, - rules: List, - rule_specs: List[ParamSpec], - residual_features: Dict[str, List[str]], - ) -> Tuple[FitResult, float]: - """Joint physical+rule fit against free-running base-sim rollouts. - - Reached when the artifact declares ``PHYSICAL_PARAM_SPECS``. Consumes - the RAW observed trajectories rather than ``base_pred_triples``: - physical parameters only manifest when momentum free-runs, which - the teacher-forced triples destroy (``State`` has no - velocities). One theta = physical + rule params, one joint fit, - so rules cannot silently absorb physics error; with no rules - this degenerates to pure identification. The identified physical - values are applied in place to the planning base env, and the - per-parameter identifiability report (posterior contraction) is - logged so null parameters are visible rather than silently - trusted. - """ - physical_specs = self._physical_param_specs - physical_names = [s.name for s in physical_specs] - # Factory, not an instance: every rollout runs in a fresh env. - fit_env = self._get_rollout_fit_env() - self._persist_fit_trajectories() - rollouts = self._rollout_fit_trajectories(residual_features) - init_params = { - s.name: s.init_value - for s in physical_specs + rule_specs - } - anchors = self.fit_prior_anchors(physical_specs) - if not rollouts: - logger.warning( - "No complete trajectories for rollout sysID; keeping the " - "declared physical-param inits unfitted.") - result = fit_params_rollout(fit_env, [], - physical_specs, - residual_features, - rules=rules, - rule_specs=rule_specs, - latent_init=self._latent_init, - anchors=anchors) - self._apply_identified_physical_params( - {n: init_params[n] - for n in physical_names}) - return result, float("nan") - - # The adjuster's third argument is the fit's own SSE-at-theta - # probe (survivor set + shared scaling - the objective the fit - # minimized), which the cross-cycle consistency check uses to - # arbitrate flagged jumps on evidence. Deliberately NOT a - # full-set SSE: trimmed (unexplainable) segments would add the - # same large error to both candidates and dilute the ratio. - outcome = run_rollout_sysid( - fit_env, - rollouts, - physical_specs, - residual_features, - rules=rules, - rule_specs=rule_specs, - latent_init=self._latent_init, - anchors=anchors, - rms_cache=self._explainability_cache, - report_adjuster=lambda result, report, sse_fn: - (self._check_cross_cycle_consistency( - result, report, physical_names, pooled_sse=sse_fn)), - held={ - **self._identified_physical_params, - **self._cycle_applied_physical - }) - if outcome.num_survivors == 0: - # NO fit ran (the result is pinned at the declared inits). - # Apply nothing: the planner keeps its standing belief - - # the previous cycle's applied values - rather than being - # reverted to baselines (or moved to this call's declared - # inits) by data that supports neither. - logger.warning( - "Rollout sysID: no explainable segments this cycle; " - "leaving the planner's physical params untouched.") - self._record_sysid_diagnostics({}, physical_names, 0, - len(rollouts), outcome.traj_rms) - return outcome.fit_result, float("nan") - inference = outcome.inference - logger.info("Identifiability (posterior/prior contraction):\n%s", - format_identifiability(inference.parameter_diagnostics)) - log_param_changes(init_params, inference.point_estimate) - self._apply_identified_physical_params(inference.selected_parameters) - self.note_carried_posterior(inference.selected_parameters, - inference.parameter_diagnostics) - if outcome.evidence is not None: - self.note_fit_evidence( - self._current_simulator_version or "harness", outcome.evidence) - # Snapshot the cycle-level decision: this (not whatever the - # agent's in-session sim.fit last applied) is what a future - # INCONSISTENT verdict holds on to. - self._cycle_applied_physical = dict(inference.selected_parameters) - # Physics-margin points for the capture gate: the fit's posterior - # widths (floored, see identifiability_report) turned into a grid - # of perturbations spanning +-1 sigma of the applied values. - self._identified_physical_sigma_points = self._physics_margin_points( - inference.selected_parameters, inference.parameter_diagnostics, - physical_specs) - if self._identified_physical_sigma_points: - logger.info("Physics-margin points for capture validation: %s", - [{k: f"{v:.4f}" - for k, v in pt.items()} - for pt in self._identified_physical_sigma_points]) - self._record_sysid_diagnostics(inference.parameter_diagnostics, - physical_names, outcome.num_survivors, - len(rollouts), outcome.traj_rms) - return outcome.fit_result, outcome.post_sse - - def _check_cross_cycle_consistency( - self, - result: FitResult, - report: Dict[str, Dict[str, Any]], - physical_names: Sequence[str], - pooled_sse: Optional[Callable[[Dict[str, float]], - float]] = None) -> None: - """Flag params whose confident MAP jumped since the previous cycle. - - The curvature probe measures local *precision*: a biased - objective yields precisely-wrong values that the probe still - stamps "identified" (observed: per-cycle friction fits 0.0585, - 0.0614, 0.0919, 0.0794, each with posterior_std ~0.003 - - mutually incompatible by many sigmas). Comparing successive - final fits in FIT space (log for log-scale params) catches - exactly this: a jump above - ``CFG.code_sim_learning_rollout_cross_cycle_sigma`` combined - sigmas sets ``Verdict.INCONSISTENT``, and the trust selection - then HOLDS the currently-applied value instead of hopping to - the new fit - neither of two mutually-incompatible confident - fits can be preferred on this evidence, and hopping churned the - belief env for whole runs (run_20260721_205821 seed1: - restitution 0.71 -> 0.52 -> 0.02 -> 0.32 -> 0.02). History - records only these final per-cycle fits, not the agent's - in-session tool fits, whose param sets churn. - - Sigma distance alone cannot tell a real correction from probe - churn, and successive cycle fits are NOT independent equals: - the new fit minimized the objective over a superset of the old - fit's data. So before holding, a flagged jump is arbitrated on - evidence via ``pooled_sse`` (the fit's own SSE-at-theta probe - over its surviving segments - the objective it minimized): - when the held value explains that data decisively worse than - the new fit - (``CFG.code_sim_learning_rollout_consistency_sse_ratio``), the - jump is accepted. Without a decisive gap the hold stands - (run_20260724_232411-style subset disagreement stays held + - hull-swept). Motivated by run_20260727_210827 seed1: a sharp - but biased 2-trajectory cycle-0 fit (0.9313, true 0.5) was held - over the 4-trajectory refit (0.4748, pooled SSE 0.14 vs ~4.4) - for the rest of the run. - """ - k = CFG.code_sim_learning_rollout_cross_cycle_sigma - fitted = result.point_estimate - scales = result.scales or ["linear"] * len(result.names) - for i, name in enumerate(result.names): - if name not in physical_names: - continue - if report.get(name, {}).get("verdict") is Verdict.ANCHORED: - # An ablation-reverted param's point estimate IS the - # baseline, not a fit. Recording it would make the next - # cycle's genuine fit read as a many-sigma jump (and - # spuriously downgrade it); keep the previous history - # entry, which holds the last real fit. - continue - post = float( - report.get(name, {}).get("posterior_std", float("nan"))) - scale = scales[i] - value = fitted[name] - prev = self._sysid_fit_history.get(name) - flagged = False - if (k > 0 and prev is not None and np.isfinite(post) - and np.isfinite(prev[1])): - prev_val, prev_std, prev_scale = prev - if prev_scale == scale: - dist = _fit_space_dist(value, prev_val, scale) - combined = float(np.sqrt(post**2 + prev_std**2)) - if combined > 0 and dist / combined > k: - n_sigma = dist / combined - pending = self._sysid_pending_fit.get(name) - if pending is not None and _fit_space_dist( - value, pending[0], scale) / max( - float(np.sqrt(post**2 + pending[1]**2)), - 1e-12) <= k: - # Two INDEPENDENT cycles agree on the new - # value: the jump was real, not probe - # overconfidence - accept it. - logger.info( - "Rollout sysID cross-cycle consistency: " - "%s jump to ~%.4f confirmed by an " - "independent refit (pending %.4f); " - "accepting the new value.", name, value, - pending[0]) - self._sysid_pending_fit.pop(name, None) - elif self._arbitrate_cross_cycle_jump( - name, fitted, prev_val, pooled_sse): - # Pooled evidence decisively prefers the - # new fit over the held value (the helper - # logs the SSE gap); accept it now instead - # of waiting a cycle for confirmation. - self._sysid_pending_fit.pop(name, None) - else: - flagged = True - logger.warning( - "Rollout sysID cross-cycle consistency: %s " - "moved %.4f -> %.4f (%.1f combined sigmas > " - "%g) since the previous cycle; the posterior " - "is overconfident.", name, prev_val, value, - n_sigma, k) - entry = report.get(name) - if (entry is not None and - entry["verdict"] is Verdict.IDENTIFIED): - entry["verdict"] = Verdict.INCONSISTENT - entry["note"] = ( - f"{prev_val:.4f} -> {value:.4f} is " - f"{n_sigma:.1f} combined sigmas; holding " - "the last trusted value, margin sweep " - "spans both") - # Both incompatible fits become hull - # candidates so the margin sweep covers - # the whole disagreement - the interval - # [0.3236, 0.6267] contained the true - # 0.5 in run_20260724_232411 seed2. - cands = set(entry.get("candidate_values", ())) - cands.update((float(prev_val), float(value))) - if pending is not None: - cands.add(float(pending[0])) - entry["candidate_values"] = sorted(cands) - self._sysid_pending_fit[name] = (value, post) - if flagged: - # Keep the trusted value as the comparison reference; - # the rejected fit waits in _sysid_pending_fit for an - # independent confirmation. Recording the rejected fit - # here would make it the NEXT cycle's reference, i.e. - # accept the hop one cycle late without any new - # evidence. - continue - self._sysid_pending_fit.pop(name, None) - self._sysid_fit_history[name] = (value, post, scale) - - @staticmethod - def _arbitrate_cross_cycle_jump( - name: str, fitted: Dict[str, float], prev_val: float, - pooled_sse: Optional[Callable[[Dict[str, float]], float]]) -> bool: - """Settle a flagged cross-cycle jump by pooled-data evidence. - - Evaluates ``pooled_sse`` (the fit's own objective over its - surviving segments) under the new joint fit and under the same - fit with ``name`` swapped back to the held value. Returns True - (accept the jump) only when the held - value's explanation is decisively worse - at least - ``CFG.code_sim_learning_rollout_consistency_sse_ratio`` times - the new fit's SSE. Anything short of decisive (including an SSE - evaluation failure) returns False and leaves the hold-and- - hull-sweep behavior in charge. - """ - ratio = CFG.code_sim_learning_rollout_consistency_sse_ratio - if pooled_sse is None or ratio <= 0: - return False - try: - sse_new = pooled_sse(dict(fitted)) - held_theta = dict(fitted) - held_theta[name] = prev_val - sse_held = pooled_sse(held_theta) - except Exception: # pylint: disable=broad-except - logger.warning( - "Rollout sysID cross-cycle arbitration: pooled SSE " - "evaluation failed for %s; holding the trusted value.", - name, - exc_info=True) - return False - if not (np.isfinite(sse_new) and np.isfinite(sse_held)): - return False - decisive = sse_held > ratio * sse_new - logger.info( - "Rollout sysID cross-cycle arbitration: %s pooled SSE %.4g at " - "the new fit %.4g vs %.4g at the held value %.4g - %s.", name, - sse_new, fitted[name], sse_held, prev_val, - ("decisively better, accepting the jump" - if decisive else "not decisive, holding")) - return decisive - - def _record_sysid_diagnostics(self, report: Dict[str, Dict[str, Any]], - physical_names: Sequence[str], - num_survivors: int, num_segments: int, - rms: List[float]) -> None: - """Digest the fit's weak spots for the next explore phase. - - Generic (domain-free) statements of what the data could not - support - unexplainable segments, parameters the rollouts do not - constrain, cross-cycle conflicts - phrased as experiment - objectives. The explorer appends this to its guidance so the - agent designs interactions that fill the gaps, instead of - relying on whatever manipulation data the tasks happen to - produce. - """ - lines: List[str] = [] - dropped = num_segments - num_survivors - if dropped > 0: - lines.append( - f"- {dropped} of {num_segments} recorded motion segments " - "were unexplainable at ANY physical parameters (best RMS " - f"{[f'{r:.3g}' for r in rms]}): their dynamics are not " - "repeatable under replay. Prefer experiments whose outcome " - "is dominated by object dynamics rather than prolonged " - "robot-object contact: actuate cleanly, then let the scene " - "evolve and settle on its own.") - for name in physical_names: - entry = report.get(name, {}) - verdict = entry.get("verdict", Verdict.UNKNOWN) - note = entry.get("note", "") - label = verdict.value + (f" ({note})" if note else "") - if verdict is Verdict.IDENTIFIED: - cands = entry.get("candidate_values", ()) - if len(cands) > 1: - lines.append( - f"- physical param '{name}': identified, but the " - "recorded segments preferred mutually-incompatible " - f"explanations spanning [{min(cands):.4g}, " - f"{max(cands):.4g}] (the physics-margin sweep " - "covers that whole hull). A clean, repeatable " - "interaction that excites this parameter and " - "little else would collapse the hull.") - continue - belief = entry.get("belief_interval") - if verdict is Verdict.WIDE and belief is not None: - anchor = entry.get("anchor") - if anchor is None: - where = "" - elif anchor < belief[0] or anchor > belief[1]: - where = (f"; the baseline {anchor:.4g} lies outside it, " - "so the data already exclude the baseline") - else: - where = f"; the baseline {anchor:.4g} lies inside it" - lines.append( - f"- physical param '{name}': the data moved it to " - f"{entry.get('map', float('nan')):.4g} but only weakly; " - f"the planner's belief is the interval [{belief[0]:.4g}, " - f"{belief[1]:.4g}]{where}. Plans are certified across " - "the whole interval, so an experiment whose observable " - "outcome DIFFERS across it would narrow the belief and " - "widen the set of certifiable plans.") - continue - interval = entry.get("flat_interval") - if (verdict in (Verdict.WEAKLY_IDENTIFIED, Verdict.NOT_IDENTIFIED) - and interval is not None and interval[0] != interval[1]): - lines.append( - f"- physical param '{name}': the data cannot " - f"distinguish values in [{interval[0]:.4g}, " - f"{interval[1]:.4g}]. An experiment whose observable " - "outcome DIFFERS across this interval would pin it " - "down.") - continue - if verdict is Verdict.ANCHORED: - # Anchor ablation handled this param correctly (the move - # was compensatory; the baseline is applied) - it is NOT - # a failed identification, so don't advise dropping it. - lines.append( - f"- physical param '{name}': the fitted move was " - "compensatory (a refit with it at its baseline explains " - "the data equally well), so the baseline was kept. An " - "experiment that excites this parameter SPECIFICALLY " - "(not jointly with the others) would distinguish the " - "two explanations.") - continue - if verdict is Verdict.INCONSISTENT: - lines.append( - f"- physical param '{name}': successive cycles produced " - f"confident but mutually-incompatible fits ({note}). " - "The objective is biased somewhere: collect a clean, " - "repeatable interaction that excites this parameter and " - "little else, so one of the two values can be refuted.") - continue - lines.append( - f"- physical param '{name}': {label}. An experiment whose " - "observable outcome CHANGES when this parameter changes " - "would identify it; if none exists, drop it from " - "PHYSICAL_PARAM_SPECS.") - self._last_sysid_diagnostics = ("\n".join(lines) if lines else "") - def _sync_tool_context(self) -> None: super()._sync_tool_context() - self._tool_context.sysid_diagnostics = (self._last_sysid_diagnostics - or None) self._tool_context.latent_tracking_available = \ self._latent_tracking_available() @@ -3073,38 +1992,6 @@ def _group_triples_by_trajectory( idx += n return groups - def _fit_parameters_recurrent( - self, - rules: List, - specs: List[ParamSpec], - base_pred_triples: List[Tuple[State, Action, State]], - residual_features: Dict[str, List[str]], - lm_seed: Optional[Tuple[np.ndarray, Optional[np.ndarray]]] = None, - ) -> Tuple[FitResult, float]: - """LM fit over the recurrent (per-trajectory) SSE. - - Counterpart to :func:`fitting.fit_rule_parameters` for rules - that carry a latent block. Re-groups the flat - ``base_pred_triples`` into per-trajectory chunks (latent threads - within a trajectory, not across) via the lengths cached in - ``self._fit_trajectories``; falls back to a single trajectory if - no grouping info exists. Delegates the actual fit/log to - :func:`fitting.fit_rule_parameters_latent` so the agent's - ``sim.fit`` surface scores latent rules through the exact - same path. - """ - groups = self._group_triples_by_trajectory(base_pred_triples) - if not groups: - logger.warning("No trajectory groups for recurrent fitting; " - "falling back to single-trajectory rollout.") - groups = [base_pred_triples] - return fit_rule_parameters_latent(rules, - specs, - groups, - self._latent_init, - residual_features, - lm_seed=lm_seed) - def _oracle_param_sse_recurrent( self, rules: List, @@ -3168,37 +2055,6 @@ def _oracle_param_sse_rollout( logger.info(" %-30s %.4f", name, val) return sse - def _attach_initial_latent(self, task: Task) -> Task: - """Seed ``task.init.latent`` with the initial latent block. - - Refinement starts at ``task.init`` (the planner's ``traj[0]``), so - the combined simulator must find a well-formed latent there. If no - ``LATENT_INIT`` was loaded (or the resulting block is empty), leave - the task alone so downstream code keeps the legacy - ``state.latent is None`` behaviour. Overrides the no-op default in - :class:`AgentModelBasedApproach`. - """ - tracker = self.make_latent_tracker() if getattr( - self, "_residual_env_cls", None) is not None else None - if tracker is not None: - return Task(init=tracker.attach(task.init, None), - goal=task.goal, - alt_goal=task.alt_goal, - goal_nl=task.goal_nl, - evaluator=task.evaluator) - if self._latent_init is None: - return task - initial_latent = init_latent(self._latent_init, self._fitted_params - or {}) - if not initial_latent: - return task - init_state = task.init.copy() - init_state.latent = initial_latent - return Task(init=init_state, - goal=task.goal, - alt_goal=task.alt_goal, - goal_nl=task.goal_nl) - def materialise_latent( self, traj: LowLevelTrajectory, @@ -3325,25 +2181,6 @@ def _compute_base_pred_triples( return [(base_env.simulate(s, a), a, s_next) for s, a, s_next in obs_triples] - @staticmethod - def _infer_residual_features_from_scan( - obs_triples: List[Tuple[State, Action, State]], - base_pred_triples: List[Tuple[State, Action, State]], - abs_tol: float = 1e-4, - rel_tol: float = 1e-3, - min_hits: int = 3, - ) -> Dict[str, List[str]]: - """Features whose base-sim prediction diverges from observation. - - Flags ``(type, feat)`` if ``|pred - obs| > rel_tol*|obs| + abs_tol`` - on at least ``min_hits`` triples. The ``min_hits`` floor keeps - one-off PyBullet jitter from leaking base-handled features into the set. - """ - del obs_triples # objects are identical across both triple lists - hits: Dict[Tuple[str, str], int] = {} - count_residual_hits(base_pred_triples, hits, abs_tol, rel_tol) - return residual_hint_from_hits(hits, min_hits) - @staticmethod def _log_feature_set_diff( a: Dict[str, List[str]], @@ -3366,23 +2203,6 @@ def _log_feature_set_diff( if only_b: logger.info(" only in %s: %s", b_label, only_b) - @staticmethod - def _format_predicate_signatures(predicates: Set[Predicate]) -> str: - """Pretty-print predicates as ``Name(type1, type2)`` lines. - - Mirrors the ``## Available Predicates`` block in - ``bilevel_sketch.build_solve_prompt``. - """ - lines = [] - for pred in sorted(predicates, key=lambda p: p.name): - type_sig = ", ".join(t.name for t in pred.types) - line = f" {pred.name}({type_sig})" - if pred.natural_language_assertion is not None: - names = [t.name for t in pred.types] - line += f" - {pred.natural_language_assertion(names)}" - lines.append(line) - return "\n".join(lines) - def _make_evaluate_trajectory_fn(self) -> Any: """Build the ``evaluate_trajectory`` helper exposed in the synthesis exec namespace (next to ``is_goal_state``). @@ -3511,72 +2331,6 @@ def _sweep_evaluation(self, evaluator: Any, states: List[State], "certified": solved == len(per_point), } - @staticmethod - def _format_trajectory_listing( - trajectories: List[LowLevelTrajectory]) -> str: - """Render a per-trajectory listing with provenance tags. - - Each interaction trajectory shows the simulator / predicates - snapshot used to generate the plan that collected it (if - tracked). Demo trajectories list as ``demo``. Listed in the same - order the agent sees them via the ``trajectories`` var. - """ - if not trajectories: - return "" - lines = ["Trajectory roster (matches the `trajectories` list):"] - for idx, traj in enumerate(trajectories): - kind = "demo" if traj.is_demo else "interaction" - try: - task_str = f"task {traj.train_task_idx}" - except AssertionError: - task_str = "task ?" - provenance: List[str] = [] - sim_v = traj.source_simulator_version - preds_v = traj.source_predicates_version - if sim_v: - provenance.append(f"sim {sim_v}") - if preds_v: - provenance.append(f"predicates {preds_v}") - tail = (f" - generated using {', '.join(provenance)}" - if provenance else "") - if traj.env_reward is not None: - solved = int( - bool(traj.env_terminated) and not traj.env_rejected) - tail += (f" - env reward={traj.env_reward:.2f} " - f"(solved={solved})") - lines.append(f" [{idx}] {kind}, {task_str}{tail}") - return "\n".join(lines) + "\n" - - def _format_objective_block(self) -> str: - """The env's public task objective (reward form), or empty. - - Emitted when a train task's evaluator states an objective. The - statement is public by design: it contains the reward FORM - (success condition + costs), never oracle quantities like the - true minimum block count. - """ - description = next( - (t.evaluator.objective_description() - for t in self._train_tasks if t.evaluator is not None - and t.evaluator.objective_description()), "") - return learn_prompts.render_objective_block(description) - - def _format_prior_state_block(self, base: str) -> str: - """Tell the agent about any simulator/predicates left over from a - previous learning cycle. - - Returns a paragraph the agent can act on (read the files first - and treat this cycle as incremental refinement) or an empty - string if no prior state exists. The base sandbox dir is scanned - for ``simulator.py`` / ``predicates.py``. - """ - prior: List[str] = [] - if os.path.isfile(os.path.join(base, "simulator.py")): - prior.append("`./simulator.py`") - if os.path.isfile(os.path.join(base, "predicates.py")): - prior.append("`./predicates.py`") - return learn_prompts.render_prior_state_block(prior) - def _simulator_load_namespace(self) -> Dict[str, Any]: """The names pre-injected when ``simulator.py`` is exec'd. @@ -3662,35 +2416,6 @@ def _load_simulator_from_module_file( # ── Static helpers ─────────────────────────────────────────── - def _write_structs_reference(self) -> str: - """Write key struct sources to the sandbox; return the agent-visible - path.""" - # pylint: disable=import-outside-toplevel,reimported - from predicators.structs import Action as _Action - from predicators.structs import LowLevelTrajectory as _LLT - from predicators.structs import Object as _Object - from predicators.structs import State as _State - from predicators.structs import Type as _Type - - source = "\n\n".join( - inspect.getsource(cls) - for cls in [_Type, _Object, _State, _Action, _LLT]) - - base = self._tool_context.sandbox_dir or self._get_log_dir() - ref_dir = os.path.join(base, "reference") - os.makedirs(ref_dir, exist_ok=True) - ref_path = os.path.join(ref_dir, "structs.py") - with open(ref_path, "w", encoding="utf-8") as f: - f.write(source) - - # Same backend-dependent agent-visible path mapping as - # _resolve_synthesis_paths. - if CFG.agent_sdk_use_local_sandbox: - return "./reference/structs.py" - if self._tool_context.sandbox_dir: - return "/sandbox/reference/structs.py" - return ref_path - def _base_sim_reference_paths(self) -> List[str]: """Agent-visible paths of the provisioned base-sim sources. @@ -3713,25 +2438,11 @@ def _base_sim_reference_paths(self) -> List[str]: type(self._base_env).__name__) return [] names = [os.path.basename(rel) for rel in src_files] - # Same backend-dependent path mapping as _write_structs_reference. + # Reference copies exist only in the local sandbox. if CFG.agent_sdk_use_local_sandbox: return [f"./reference/base_sim/{n}" for n in names] - if CFG.agent_sdk_use_docker_sandbox: - return [f"/sandbox/reference/base_sim/{n}" for n in names] return [] - @staticmethod - def _extract_obs_triples( - trajectories: List[LowLevelTrajectory], - ) -> List[Tuple[State, Action, State]]: - """Extract observed (s_t, action_t, s_{t+1}) triples.""" - triples: List[Tuple[State, Action, State]] = [] - for traj in trajectories: - for i in range(len(traj.actions)): - triples.append( - (traj.states[i], traj.actions[i], traj.states[i + 1])) - return triples - def _recreate_base_env(self) -> None: """Reconnect after a PyBullet physics-server crash.""" self._rebuild_base_env("PyBullet physics client crashed; recreating " @@ -4088,50 +2799,6 @@ def step(state: State, action: Action) -> State: return make_stepper - def _build_synthesis_system_prompt(self) -> str: - """Compose the synthesis system prompt from the templates. - - Per-instance choices: the rule signature (flag-gated on - ``CFG.partially_observable``, which also swaps the env's - observation and the GT simulator module, so prompt and world - never disagree; under the flag only the recurrent 5-arg form is - shown), the optional PHYSICAL_PARAM_SPECS section (env parameter - menu), the scene-visualization hint, and the subclass extras. - """ - return learn_prompts.build_learn_system_prompt( - partially_observable=CFG.partially_observable, - residual_rule_signature=self._residual_rule_signature(), - scene_viz_hint=self._scene_viz_hint(), - physical_params_section=self._physical_params_prompt_section(), - extra_sections=self._extra_synthesis_system_prompt_sections(), - latent_extra_sections=self._extra_synthesis_latent_sections(), - workflow_extra=self._synthesis_workflow_extra(), - declared_params_only=CFG.agent_sim_learn_declared_params_only, - ) - - def _synthesis_workflow_extra(self) -> str: - """Extra text appended to the Workflow list. - - Subclasses with additional deliverables (e.g. invented - predicates) extend the workflow here so the numbered list stays - the single authoritative loop description. - """ - return "" - - @staticmethod - def _scene_viz_hint() -> str: - """The find-the-anchor-offset sentence. - - The probe is unconditional in synthesis sessions, so the hint - always names its staging + overlay surface. - """ - return ("use the `sim` probe in `run_python`: " - "`sim.reset(task_idx=..., " - "mods={...})` to stage a representative state from each " - "bucket and `sim.render(label, annotations=[...])` to " - "overlay, on one render, the recorded origin and the " - "positions where the effect did vs. did not fire") - def _physical_params_prompt_section(self) -> str: """Markdown for the optional PHYSICAL_PARAM_SPECS (system-ID) block. @@ -4148,17 +2815,3 @@ def _physical_params_prompt_section(self) -> str: if callable(getter): info = getter() or {} return learn_prompts.render_physical_params_section(info) - - def _residual_rule_signature(self) -> str: - """The ``def`` line used in the geometric-gate example. - - Matches the canonical rule signature the prompt advertises - (``CFG.partially_observable`` selects it) so the worked example - doesn't contradict it. - """ - if CFG.partially_observable: - # The example's body reads `state`; bind it so the recurrent - # form is a runnable rule, not a signature over a foreign name. - return ("def residual_rule(observation, latent, history, " - "updates, params):\n state = observation") - return "def residual_rule(state, updates, params):" diff --git a/predicators/approaches/agent_sim_predicate_invention_approach.py b/predicators/approaches/agent_sim_predicate_invention_approach.py index 4e07a04cb3..b1b6a3cde7 100644 --- a/predicators/approaches/agent_sim_predicate_invention_approach.py +++ b/predicators/approaches/agent_sim_predicate_invention_approach.py @@ -41,7 +41,6 @@ import os from typing import Any, Dict, FrozenSet, List, Set, Tuple -from predicators.agent_sdk import learn_prompts from predicators.agent_sdk.tools import _SnapshotTarget, \ finalize_versioned_snapshot, make_predicate_quality_loader from predicators.approaches.agent_sim_learning_approach import \ @@ -162,61 +161,6 @@ def _build_write_snapshot_targets( )) return targets - def _extra_synthesis_message(self, extra_paths: Dict[str, str]) -> str: - message = learn_prompts.render_predicate_invention_message( - extra_paths["predicates_file_for_agent"], - self._format_goal_nl_block()) - return message + self._chained_extra_message(extra_paths) - - def _chained_extra_message(self, extra_paths: Dict[str, str]) -> str: - """The base class's extra message (the partial-observability note under - ``CFG.partially_observable``), separated for appending.""" - base = super()._extra_synthesis_message(extra_paths) - return "\n\n" + base if base else "" - - def _format_goal_nl_block(self) -> str: - """Render the deduped natural-language goals for the train tasks. - - Returns an empty string only if every task is missing a - ``goal_nl``, but ``__init__`` asserts they're present, so in - practice this always returns a non-empty block. - """ - seen: List[str] = [] - for task in self._train_tasks: - nl = task.goal_nl - if nl and nl not in seen: - seen.append(nl) - if not seen: - return "" - if len(seen) == 1: - return f"Goal (natural language): {seen[0]}\n\n" - bullets = "\n".join(f" - {g}" for g in seen) - return f"Goals across train tasks (natural language):\n{bullets}\n\n" - - def _synthesis_workflow_extra(self) -> str: - # The base workflow's validation step depends on invented - # predicates: sketches can only reference predicates that exist. - return learn_prompts.render_predicate_workflow_extra() - - def _extra_synthesis_system_prompt_sections(self) -> List[str]: - # The scene workbench is the sim probe inside run_python (the - # probe is unconditional in synthesis sessions). - workbench = ("the `sim` probe in `run_python` as scene workbench " - "(`sim.reset(task_idx=..., mods={...})` to stage " - "states, `sim.render(label, annotations=[...])` " - "to render with overlays)") - sections = super()._extra_synthesis_system_prompt_sections() - sections.append( - learn_prompts.render_predicate_invention_section(workbench)) - return sections - - def _extra_synthesis_latent_sections(self) -> List[str]: - # The predicate-side latent guidance belongs to invention arms - # only and follows the simulator-side tutorial it refers to. - sections = super()._extra_synthesis_latent_sections() - sections.append(learn_prompts.render_predicate_latent_section()) - return sections - def _post_synthesis_loading( self, extra_paths: Dict[str, str], diff --git a/predicators/approaches/continual_play_mixin.py b/predicators/approaches/continual_play_mixin.py index 8fd77bca64..ecff81e57f 100644 --- a/predicators/approaches/continual_play_mixin.py +++ b/predicators/approaches/continual_play_mixin.py @@ -28,17 +28,15 @@ A reset starts a new episode on the same level without ending the round. A level can span several rounds if the agent returns before settling it. -Why a mixin. The arms' learning and session machinery live in the -phased approach classes (``AgentModelFreeApproach`` and its +Why a mixin. The arms' model and session machinery live in the +approach classes below them (``AgentModelFreeApproach`` and its ``AgentSimPredicateInventionApproach`` descendant), where the simulator -synthesis, the parameter fit, predicate invention, the sandbox and the +loading, the parameter belief, predicate invention, the sandbox and the session managers are implemented. An arm keeps that class as its base -and mixes this loop in front of it, the way ``AgentSessionMixin`` and -``SamplerLearningMixin`` add their concerns; the phased loop's own entry -points (``_solve``, the explorers, the learning sessions) are simply -unused under the protocol. The mixin has no base class of its own, so -there is no diamond, and what it needs from its host is declared below -as the host contract. +and mixes this loop in front of it, the way ``AgentSessionMixin`` adds +its concerns. The mixin has no base class of its own, so there is no +diamond, and what it needs from its host is declared below as the host +contract. The harness never chooses for the agent: whether to act, reset, model or give up is decided inside the conversation; the loop only services @@ -542,8 +540,6 @@ def _episodes_to_trajectories(self, episodes: Sequence[Dict[str, Any]], self, "_current_simulator_version", None), _source_predicates_version=getattr( self, "_current_predicates_version", None), - _source_samplers_version=getattr( - self, "_current_samplers_version", None), _env_reward=ep.get("reward"), _env_terminated=ep.get("terminated"), )) diff --git a/predicators/approaches/gnn_dynamics_approach.py b/predicators/approaches/gnn_dynamics_approach.py deleted file mode 100644 index 84cc24bd4e..0000000000 --- a/predicators/approaches/gnn_dynamics_approach.py +++ /dev/null @@ -1,459 +0,0 @@ -"""A neural dynamics baseline (paper arm C5): a GNN transition model over the -object-centric state features, trained on option-level transitions from the -same interaction data the agent arms collect, and planned against by random -shooting through the learned model. - -The model predicts, for a state and a ground option, the change of every -object feature and the number of low-level steps the option takes. The -last ``gnn_dynamics_history_len`` pre-option states ride along as node -features so a mechanism that is not visible in one observation (a cure -that started a few options ago) can be inferred from the recent past. -Planning samples random applicable options, rolls them through the -model, and accepts the first sequence whose predicted final state -satisfies the goal under the env's (oracle) predicates; by default the -plan is re-shot from the observed state every time an option terminates. -""" - -import logging -import time -from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple - -import dill as pkl -import numpy as np -import torch -import torch.nn -import torch.optim -from gym.spaces import Box -from torch.utils.data import DataLoader - -from predicators import utils -from predicators.approaches import ApproachFailure, ApproachTimeout, \ - BaseApproach -from predicators.explorers import create_explorer -from predicators.gnn.gnn import EncodeProcessDecode, setup_graph_net -from predicators.gnn.gnn_utils import GraphDictDataset, compute_normalizers, \ - get_single_model_prediction, graph_batch_collate, normalize_graph, \ - train_model -from predicators.nsrt_learning.segmentation import segment_trajectory -from predicators.settings import CFG -from predicators.structs import Action, Dataset, DummyOption, \ - InteractionRequest, InteractionResult, LowLevelTrajectory, \ - ParameterizedOption, Predicate, State, Task, Type, _Option - -# One training example: the recent pre-option states (oldest first), the -# pre-option state, the ground option, the post-option state, and the -# option's low-level step count. -_Example = Tuple[List[State], State, _Option, State, int] - - -class GNNDynamicsShootingApproach(BaseApproach): - """GNN transition model over object features + shooting planner.""" - - def __init__(self, initial_predicates: Set[Predicate], - initial_options: Set[ParameterizedOption], types: Set[Type], - action_space: Box, train_tasks: List[Task]) -> None: - super().__init__(initial_predicates, initial_options, types, - action_space, train_tasks) - self._sorted_options = sorted(self._initial_options, - key=lambda o: o.name) - self._trajectories: List[LowLevelTrajectory] = [] - self._online_learning_cycle = 0 - self._requests_train_task_idxs: List[int] = [] - self._gnn: Optional[EncodeProcessDecode] = None - self._data_exemplar: Optional[Tuple[Dict, Dict]] = None - self._input_normalizers: Optional[Dict] = None - self._target_normalizers: Optional[Dict] = None - # Node feature layout: type one-hot, current features, history - # features (per lag), option-argument slot one-hot. - self._type_to_index: Dict[str, int] = {} - self._feat_to_index: Dict[str, int] = {} - self._max_option_objects = 0 - self._max_option_params = 0 - self._mse_loss = torch.nn.MSELoss() - - @classmethod - def get_name(cls) -> str: - return "gnn_dynamics_shooting" - - @property - def is_learning_based(self) -> bool: - return True - - # ── Data ───────────────────────────────────────────────────── - - def _generate_examples(self) -> List[_Example]: - """Option-level transitions from every stored trajectory.""" - examples: List[_Example] = [] - history_len = CFG.gnn_dynamics_history_len - for traj in self._trajectories: - if not traj.actions or not traj.actions[0].has_option(): - continue - segments = segment_trajectory(traj, self._initial_predicates) - pre_states: List[State] = [] - for segment in segments: - state = segment.states[0] - history = pre_states[-history_len:] if history_len else [] - examples.append((list(history), state, segment.get_option(), - segment.states[-1], len(segment.actions))) - pre_states.append(state) - return examples - - def _setup_fields(self, examples: Sequence[_Example]) -> None: - types: Set[str] = set() - feats: Set[str] = set() - max_objects = 0 - max_params = 0 - for _, state, option, _, _ in examples: - for obj in state: - types.add(obj.type.name) - feats.update(obj.type.feature_names) - max_objects = max(max_objects, len(option.objects)) - max_params = max(max_params, option.params.shape[0]) - # Every option's argument count and parameter box, not just the - # ones seen: a test-time sample of an unseen option must graphify. - for param_opt in self._sorted_options: - max_objects = max(max_objects, len(param_opt.types)) - max_params = max(max_params, param_opt.params_space.shape[0]) - self._type_to_index = {t: i for i, t in enumerate(sorted(types))} - self._feat_to_index = {f: i for i, f in enumerate(sorted(feats))} - self._max_option_objects = max_objects - self._max_option_params = max_params - - def _feature_row(self, state: State, obj: Any) -> np.ndarray: - row = np.zeros(len(self._feat_to_index)) - for feat, val in zip(obj.type.feature_names, state[obj]): - row[self._feat_to_index[feat]] = val - return row - - def _graphify_input(self, history: Sequence[State], state: State, - option: _Option) -> Tuple[Dict, Dict[Any, int]]: - objects = list(state) - object_to_node = {o: i for i, o in enumerate(objects)} - n_obj = len(objects) - n_types = len(self._type_to_index) - n_feats = len(self._feat_to_index) - history_len = CFG.gnn_dynamics_history_len - n_node_feats = (n_types + n_feats + history_len * n_feats + - self._max_option_objects) - nodes = np.zeros((n_obj, n_node_feats)) - arg_slots = {o: i for i, o in enumerate(option.objects)} - for obj in objects: - i = object_to_node[obj] - nodes[i, self._type_to_index[obj.type.name]] = 1 - current = self._feature_row(state, obj) - nodes[i, n_types:n_types + n_feats] = current - # History as differences to the current features, most - # recent lag first; missing lags stay zero (no change). - for lag in range(history_len): - if lag < len(history): - past = history[-1 - lag] - if obj in past.data: - past_row = self._feature_row(past, obj) - start = n_types + n_feats * (1 + lag) - nodes[i, start:start + n_feats] = past_row - current - if obj in arg_slots: - nodes[i, n_types + n_feats * (1 + history_len) + - arg_slots[obj]] = 1 - # Fully connected (no self loops) so effects can flow between - # objects; edge features: constant, both endpoints are arguments. - senders, receivers, edges = [], [], [] - for s in range(n_obj): - for r in range(n_obj): - if s == r: - continue - senders.append(s) - receivers.append(r) - both = float(objects[s] in arg_slots - and objects[r] in arg_slots) - edges.append([1.0, both]) - n_edge = len(edges) - onehot = np.zeros(len(self._sorted_options)) - onehot[self._sorted_options.index(option.parent)] = 1 - params = np.zeros(self._max_option_params) - params[:option.params.shape[0]] = option.params - graph = { - "n_node": np.array(n_obj), - "nodes": nodes, - "n_edge": np.reshape(n_edge, [1]).astype(np.int64), - "edges": np.reshape(edges, [n_edge, 2]), - "senders": np.reshape(senders, [n_edge]).astype(np.int64), - "receivers": np.reshape(receivers, [n_edge]).astype(np.int64), - "globals": np.r_[onehot, params], - } - return graph, object_to_node - - def _graphify_target(self, state: State, next_state: State, - num_actions: int, graph_input: Dict, - object_to_node: Dict[Any, int]) -> Dict: - n_obj = len(object_to_node) - nodes = np.zeros((n_obj, len(self._feat_to_index))) - for obj, i in object_to_node.items(): - nodes[i] = (self._feature_row(next_state, obj) - - self._feature_row(state, obj)) - return { - "n_node": graph_input["n_node"], - "nodes": nodes, - "n_edge": graph_input["n_edge"], - "edges": np.zeros((int(graph_input["n_edge"][0]), 1)), - "senders": graph_input["senders"], - "receivers": graph_input["receivers"], - "globals": np.array([num_actions / float(max(CFG.horizon, 1))]), - } - - # ── Learning ───────────────────────────────────────────────── - - def learn_from_offline_dataset(self, dataset: Dataset) -> None: - self._trajectories = list(dataset.trajectories) - self._learn_model() - self._save(None) - - def get_interaction_requests(self) -> List[InteractionRequest]: - explorer = create_explorer(CFG.explorer, self._initial_predicates, - self._initial_options, self._types, - self._action_space, self._train_tasks) - requests: List[InteractionRequest] = [] - self._requests_train_task_idxs = [] - for _ in range(CFG.online_nsrt_learning_requests_per_cycle): - task_idx = int(self._rng.choice(len(self._train_tasks))) - policy, termination_fn = explorer.get_exploration_strategy( - task_idx, CFG.timeout) - requests.append( - InteractionRequest(train_task_idx=task_idx, - act_policy=policy, - query_policy=lambda s: None, - termination_function=termination_fn)) - self._requests_train_task_idxs.append(task_idx) - return requests - - def restore_interaction_requests(self, train_task_idxs: List[int]) -> None: - self._requests_train_task_idxs = list(train_task_idxs) - - def learn_from_interaction_results( - self, results: Sequence[InteractionResult]) -> None: - assert len(results) == len(self._requests_train_task_idxs) - for task_idx, result in zip(self._requests_train_task_idxs, results): - self._trajectories.append( - LowLevelTrajectory(result.states, - result.actions, - _is_demo=False, - _train_task_idx=task_idx)) - self._learn_model() - self._save(self._online_learning_cycle) - self._online_learning_cycle += 1 - - def _learn_model(self) -> None: - examples = self._generate_examples() - if not examples: - logging.warning("GNN dynamics: no option-level transitions yet; " - "keeping the previous model.") - return - self._setup_fields(examples) - graph_inputs, graph_targets = [], [] - for history, state, option, next_state, num_actions in examples: - graph_input, object_to_node = self._graphify_input( - history, state, option) - graph_inputs.append(graph_input) - graph_targets.append( - self._graphify_target(state, next_state, num_actions, - graph_input, object_to_node)) - self._data_exemplar = (graph_inputs[0], graph_targets[0]) - self._gnn = setup_graph_net(GraphDictDataset([graph_inputs[0]], - [graph_targets[0]]), - num_steps=CFG.gnn_num_message_passing, - layer_size=CFG.gnn_layer_size) - if CFG.gnn_do_normalization: - self._input_normalizers = compute_normalizers(graph_inputs) - self._target_normalizers = compute_normalizers(graph_targets) - graph_inputs = [ - normalize_graph(g, self._input_normalizers) - for g in graph_inputs - ] - graph_targets = [ - normalize_graph(g, self._target_normalizers) - for g in graph_targets - ] - num_validation = (max(1, int(len(examples) * 0.1)) - if CFG.gnn_use_validation_set else 0) - train_set = GraphDictDataset(graph_inputs[num_validation:], - graph_targets[num_validation:]) - val_set = GraphDictDataset(graph_inputs[:num_validation], - graph_targets[:num_validation]) - dataloaders = { - "train": - DataLoader(train_set, - batch_size=CFG.gnn_batch_size, - shuffle=False, - num_workers=0, - collate_fn=graph_batch_collate), - "val": - DataLoader(val_set, - batch_size=CFG.gnn_batch_size, - shuffle=False, - num_workers=0, - collate_fn=graph_batch_collate), - } - optimizer = torch.optim.Adam(self._gnn.parameters(), - lr=CFG.gnn_learning_rate, - weight_decay=CFG.gnn_weight_decay) - logging.info( - "Training GNN dynamics on %d option transitions from " - "%d trajectories.", len(examples), len(self._trajectories)) - best = train_model(self._gnn, - dataloaders, - optimizer=optimizer, - criterion=self._mse_loss, - global_criterion=self._mse_loss, - num_epochs=CFG.gnn_num_epochs, - do_validation=CFG.gnn_use_validation_set) - self._gnn.load_state_dict(best) - - # ── Checkpointing ──────────────────────────────────────────── - - def _save(self, online_learning_cycle: Optional[int]) -> None: - info = { - "trajectories": - self._trajectories, - "online_learning_cycle": - self._online_learning_cycle, - "exemplar": - self._data_exemplar, - "state_dict": - (self._gnn.state_dict() if self._gnn is not None else None), - "type_to_index": - self._type_to_index, - "feat_to_index": - self._feat_to_index, - "max_option_objects": - self._max_option_objects, - "max_option_params": - self._max_option_params, - "input_normalizers": - self._input_normalizers, - "target_normalizers": - self._target_normalizers, - } - path = (f"{utils.get_approach_save_path_str()}_" - f"{online_learning_cycle}.gnn") - with open(path, "wb") as f: - pkl.dump(info, f) - - def load(self, online_learning_cycle: Optional[int]) -> None: - path = (f"{utils.get_approach_load_path_str()}_" - f"{online_learning_cycle}.gnn") - with open(path, "rb") as f: - info = pkl.load(f) - self._trajectories = info["trajectories"] - self._online_learning_cycle = info["online_learning_cycle"] - self._data_exemplar = info["exemplar"] - self._type_to_index = info["type_to_index"] - self._feat_to_index = info["feat_to_index"] - self._max_option_objects = info["max_option_objects"] - self._max_option_params = info["max_option_params"] - self._input_normalizers = info["input_normalizers"] - self._target_normalizers = info["target_normalizers"] - if info["state_dict"] is not None and self._data_exemplar is not None: - example_input, example_target = self._data_exemplar - self._gnn = setup_graph_net(GraphDictDataset([example_input], - [example_target]), - num_steps=CFG.gnn_num_message_passing, - layer_size=CFG.gnn_layer_size) - self._gnn.load_state_dict(info["state_dict"]) - - # ── Prediction ─────────────────────────────────────────────── - - def predict_next_state(self, history: Sequence[State], state: State, - option: _Option) -> Tuple[State, int]: - """The model's post-option state and low-level step count.""" - assert self._gnn is not None, "Learn a model before predicting." - graph_input, object_to_node = self._graphify_input( - history, state, option) - if CFG.gnn_do_normalization: - assert self._input_normalizers is not None - graph_input = normalize_graph(graph_input, self._input_normalizers) - out = get_single_model_prediction(self._gnn, graph_input) - if CFG.gnn_do_normalization: - assert self._target_normalizers is not None - out = normalize_graph(out, self._target_normalizers, invert=True) - next_state = state.copy() - for obj, i in object_to_node.items(): - delta = out["nodes"][i] - for feat in obj.type.feature_names: - next_state.set( - obj, feat, - state.get(obj, feat) + - float(delta[self._feat_to_index[feat]])) - num_actions = max( - 1, int(round(float(out["globals"][0]) * max(CFG.horizon, 1)))) - return next_state, num_actions - - # ── Planning ───────────────────────────────────────────────── - - def _shoot(self, task: Task, init_state: State, history: Sequence[State], - deadline: float) -> Optional[List[_Option]]: - """Random shooting through the learned model. - - Returns the first sampled option sequence whose predicted final - state satisfies the goal, or None when the tries or the deadline - run out. - """ - for _ in range(CFG.gnn_dynamics_shooting_max_tries): - if time.perf_counter() > deadline: - return None - state = init_state - past = list(history) - plan: List[_Option] = [] - num_actions = 0 - for _ in range(CFG.gnn_dynamics_max_plan_length): - if task.goal_holds(state): - return plan - option = utils.sample_applicable_option( - self._sorted_options, state, self._rng) - if option is None: - break - next_state, k = self.predict_next_state(past, state, option) - plan.append(option) - past.append(state) - state = next_state - num_actions += k - if num_actions > CFG.horizon: - break - if task.goal_holds(state) and plan: - return plan - return None - - def _solve(self, task: Task, timeout: int) -> Callable[[State], Action]: - if self._gnn is None: - raise ApproachFailure("GNN dynamics model has not been learned.") - deadline = time.perf_counter() + timeout - observed: List[State] = [] - plan: List[_Option] = [] - cur_option: _Option = DummyOption - replan = CFG.gnn_dynamics_replan_every_option - - def _next_option(state: State) -> _Option: - nonlocal plan - if replan or not plan: - shot = self._shoot(task, state, observed, deadline) - if shot is None: - if time.perf_counter() > deadline: - raise ApproachTimeout( - "GNN dynamics shooting timed out.") - raise ApproachFailure( - "GNN dynamics shooting found no plan that reaches " - "the goal under the learned model.") - plan = list(shot) - option = plan.pop(0) - if not option.initiable(state): - raise ApproachFailure( - f"Planned option {option.name} is not initiable in the " - "observed state.") - return option - - def _policy(state: State) -> Action: - nonlocal cur_option - if cur_option is DummyOption or cur_option.terminal(state): - cur_option = _next_option(state) - observed.append(state) - return cur_option.policy(state) - - return _policy diff --git a/predicators/approaches/sampler_learning_mixin.py b/predicators/approaches/sampler_learning_mixin.py deleted file mode 100644 index 023ba2c429..0000000000 --- a/predicators/approaches/sampler_learning_mixin.py +++ /dev/null @@ -1,431 +0,0 @@ -"""Parameterized (per-skill) sampler learning for the sim-learning approach. - -A parameterized sampler is keyed by option name and authored once; the -ground level of the sampler hierarchy (per-step ``GroundSampler`` from a -sketch ``~ [widths]`` region annotation) is not learned and lives in -``bilevel_sketch``, overriding the parameterized sampler per step. - -Samplers are a first-class artifact of the base sim-learning approach -(gated by ``CFG.agent_sim_learn_parameterized_samplers``), not a subclass -extension like predicates — so they are woven into -``AgentSimLearningApproach._synthesize_with_agent`` and -``_learn_simulator`` directly rather than via the ``_extra_synthesis_*`` -hooks, which keeps them independent of the predicate subclass's -(non-super-calling) hook overrides. When a sim-synthesis session runs -(``oracle_sim_program=False``) the sampler tool/snapshot/message ride -along in it; when none runs (``oracle_sim_program=True``) they get a -dedicated session via :meth:`_synthesize_samplers_standalone`. - -This mixin owns everything sampler-specific: mode resolution (learn vs. -ground truth), sandbox path bindings, the synthesis tool/snapshot/message -builders, loading ``LEARNED_SAMPLERS`` from file, and the standalone -synthesis session. The host approach keeps only the call sites. -""" -import logging -import os -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, cast - -from predicators.agent_sdk.session_base import AgentSessionFatalError, \ - query_fatal_error -from predicators.agent_sdk.tools import _SnapshotTarget, \ - create_synthesis_tools, finalize_versioned_snapshot, make_sampler_loader -from predicators.agent_sdk.tools.digests import render_options_digest -from predicators.code_sim_learning.fit_space import ParamSpec -from predicators.ground_truth_models import get_gt_samplers -from predicators.settings import CFG -from predicators.structs import Action, LowLevelTrajectory, \ - ParameterizedOption, ParameterizedSampler, Predicate, State, Task, Type - -if TYPE_CHECKING: - from predicators.agent_sdk.synthesis_backend import SynthesisBackend - from predicators.agent_sdk.tools import ToolContext - -logger = logging.getLogger(__name__) - - -class SamplerLearningMixin: - """Per-skill sampler synthesis, loading, and oracle installation. - - Mixed into :class:`AgentSimLearningApproach`. Holds the - sampler-learning state (``_do_synthesize_samplers``, - ``_current_samplers_version``) — the host ``__init__`` must call - :meth:`_init_sampler_learning_state`. - """ - - # ── Host-class contract ───────────────────────────────────── - # Everything below is provided by the host approach (its - # AgentSessionMixin / BaseApproach ancestry or the host class - # itself). Declared under TYPE_CHECKING only, so these never - # shadow the real implementations in the MRO at runtime. - if TYPE_CHECKING: - _tool_context: "ToolContext" - _train_tasks: List[Task] - _types: Set[Type] - _fitted_params: Dict[str, float] - _learning_mode: bool - _synthesized_samplers: Dict[str, ParameterizedSampler] - - def _learning_cycle_index(self) -> int: - raise NotImplementedError - - def _get_log_dir(self) -> str: - raise NotImplementedError - - def _get_all_predicates(self) -> Set[Predicate]: - raise NotImplementedError - - def _get_all_options(self) -> Set[ParameterizedOption]: - raise NotImplementedError - - def _get_synthesis_tool_names(self) -> Optional[List[str]]: - raise NotImplementedError - - def _build_synthesis_exec_ns( - self, - trajectories: List[LowLevelTrajectory]) -> Dict[str, Any]: - raise NotImplementedError - - def _query_agent_sync(self, message: str, - **query_kwargs: Any) -> List[Dict[str, Any]]: - raise NotImplementedError - - def _ensure_agent_session(self) -> None: - raise NotImplementedError - - def _close_agent_session(self) -> None: - raise NotImplementedError - - @staticmethod - def _build_synthesis_session_hooks( - targets: List[_SnapshotTarget], - sandbox_dir: str) -> Dict[str, list]: - raise NotImplementedError - - @staticmethod - def _format_predicate_signatures(predicates: Set[Predicate]) -> str: - raise NotImplementedError - - def _init_sampler_learning_state(self) -> None: - """Initialize sampler state; called from the host ``__init__``.""" - # Snapshot tag of the most recent samplers file committed by the - # synthesis agent — used to stamp newly collected online - # trajectories with their source-version provenance. - self._current_samplers_version: Optional[str] = None - # Whether this run learns samplers (vs. using ground-truth ones). - # Refined per cycle in _learn_simulator once GT availability is - # known; this default is what the synthesis-session tool surface - # reads. - self._do_synthesize_samplers: bool = ( - CFG.agent_sim_learn_parameterized_samplers - and not CFG.agent_sim_learn_oracle_samplers) - - @staticmethod - def _samplers_enabled() -> bool: - """Whether per-skill samplers are used at all this run.""" - return CFG.agent_sim_learn_parameterized_samplers - - def _maybe_install_oracle_samplers(self) -> None: - """Resolve sampler mode for this cycle and install GT ones if used. - - Sets ``self._do_synthesize_samplers`` (learn vs. use ground - truth). When ``agent_sim_learn_oracle_samplers`` is on and the - env provides ground-truth samplers, installs them and skips - synthesis; if none exist, warns and falls back to synthesis. - """ - gt_samplers = None - if self._samplers_enabled() and CFG.agent_sim_learn_oracle_samplers: - gt_samplers = get_gt_samplers(CFG.env) - if gt_samplers: - self._synthesized_samplers = dict(gt_samplers) - self._current_samplers_version = "oracle" - logger.info("Using %d ground-truth sampler(s): %s", - len(gt_samplers), ", ".join(sorted(gt_samplers))) - else: - logger.warning( - "agent_sim_learn_oracle_samplers=True but no ground-truth " - "samplers for env %s; falling back to synthesis.", CFG.env) - self._do_synthesize_samplers = (self._samplers_enabled() - and not gt_samplers) - - def _sampler_paths(self, base: str) -> Dict[str, str]: - """Sandbox path bindings for samplers.py (host + agent-visible).""" - samplers_file = os.path.join(base, "samplers.py") - samplers_versions_dir = os.path.join(base, "samplers_versions") - if CFG.agent_sdk_use_local_sandbox: - samplers_file_for_agent = "./samplers.py" - elif self._tool_context.sandbox_dir: - samplers_file_for_agent = "/sandbox/samplers.py" - else: - samplers_file_for_agent = samplers_file - return { - "samplers_file": samplers_file, - "samplers_versions_dir": samplers_versions_dir, - "samplers_file_for_agent": samplers_file_for_agent, - } - - def _install_sampler_surface(self, paths: Dict[str, str]) -> None: - """Register the ``sim.samplers()`` loader for a synthesis session.""" - self._tool_context.probe_artifact_loaders["samplers"] = \ - make_sampler_loader( - samplers_file=paths["samplers_file"], - samplers_versions_dir=paths["samplers_versions_dir"], - approach=self, - cycle_index_provider=self._learning_cycle_index, - ) - - def _sampler_snapshot_target(self, paths: Dict[str, - str]) -> _SnapshotTarget: - """Snapshot target that versions samplers.py on every Write/Edit.""" - return _SnapshotTarget( - live_file=paths["samplers_file"], - versions_dir=paths["samplers_versions_dir"], - artifact_name="samplers", - cycle_index_provider=self._learning_cycle_index, - ) - - def _sampler_synthesis_message(self, paths: Dict[str, str]) -> str: - """Instructions appended to the agent's first synthesis message.""" - path = paths["samplers_file_for_agent"] - # The ground channel exists only when its flag is on; do not - # describe it to sessions that cannot use it. - ground_note = "" - if CFG.agent_bilevel_ground_samplers: - ground_note = ( - "\nSamplers here are the reusable cross-task prior: " - "refinement uses yours on every draw of that option, in " - "every sketch and every task. A sketch step that carries " - "its own `~` ground-sampler annotation (a `~ [widths]` " - "window or `~ name` from ground_samplers.py) bypasses " - "yours for that step (precedence: ground sampler > " - "parameterized sampler > uniform).") - return f"""\ -## Per-Skill Sampler Synthesis - -Backtracking refinement draws each option's continuous parameters \ -*uniformly* from its params box by default. When a sketch step's subgoal \ -pins the parameters into a tiny region (e.g. a placement that must land \ -within a few cm of an exact point and at a specific orientation), uniform \ -sampling almost never hits it and refinement exhausts its budget. Fix this \ -by writing per-skill samplers to `{path}` as a dict \ -`LEARNED_SAMPLERS = {{"OptionName": sampler_fn, ...}}` keyed by option name. - -Each sampler has signature \ -`fn(state, subgoal_atoms, rng, objects) -> params` (the same signature as \ -the env's NSRT samplers) where: -- `state` is the current `State` (read object features with `state.get(obj, "feat")`), -- `subgoal_atoms` is the set of `GroundAtom`s the step must establish — \ -read the target relation here (e.g. an `InFront`/at-target atom names the \ -two objects whose geometry the placement must satisfy) and compute the \ -parameters that achieve it. At steps with NO subgoal annotation this set \ -is EMPTY — the sampler must not crash on `set()`; fall back to a default \ -or uniform draw, -- `rng` is a `numpy` `Generator` (use it for small jitter so retries differ), -- `objects` is the list of typed objects bound to this option call. -Return a `float32` array whose length matches the option's params box \ -(see the Options digest in your prompt for the dimension and ranges); \ -refinement clips it to that box, so stay within the ranges. -{ground_note} - -Aim the parameters at the subgoal geometrically (then add a little `rng` \ -jitter); do NOT just return uniform draws. Read the option signatures \ -from the Options digest in your prompt and the predicate classifiers \ -(for the subgoal geometry) with the predicate listing above. - -Workflow: write `{path}`, call `sim.samplers()` (snapshots + installs \ -them and sanity-checks shape/box), then call `sim.refine` \ -with a sketch using those options — the samples-to-refine count should \ -drop sharply versus uniform. Iterate with `Edit` and re-run. Every \ -successful Write/Edit of `{path}` is snapshotted to `samplers_versions/` \ -as `cycle_XXX_vers_YYY_samplers.py`.""" - - def _finalize_and_load_samplers(self, paths: Dict[str, str]) -> None: - """Snapshot the final samplers.py and load it into approach state.""" - tag = finalize_versioned_snapshot( - paths["samplers_file"], - paths["samplers_versions_dir"], - cycle_idx=self._learning_cycle_index(), - artifact_name="samplers", - ) - if tag is not None: - self._current_samplers_version = tag - logger.info("Final samplers snapshot: %s", tag) - loaded = self._load_samplers_from_module_file(paths["samplers_file"]) - self._synthesized_samplers = loaded - logger.info("Loaded %d per-skill sampler(s) from %s.", len(loaded), - paths["samplers_file"]) - for name in sorted(loaded): - logger.info(" sampler: %s", name) - - def _load_samplers_from_module_file( - self, path: str) -> Dict[str, ParameterizedSampler]: - """Load LEARNED_SAMPLERS from ``path``; validate each entry. - - Mirrors ``_load_predicates_from_module_file``. Returns an empty - dict on missing file or exec failure (samplers are optional). - Validation (unknown option names, non-callables) is shared with - ``sim.samplers()`` via ``load_learned_samplers``. - """ - # pylint: disable=import-outside-toplevel - from predicators.agent_sdk.proposal_exec import build_exec_context, \ - load_learned_samplers - from predicators.agent_sdk.tools import _ParamsView - - # pylint: enable=import-outside-toplevel - # ParamSpec is imported at module scope (used by exec'd samplers - # that close over learned params, mirroring the predicate loader). - - if not os.path.isfile(path): - logger.info("No samplers file at %s; sampler set is empty.", path) - return {} - - with open(path, "r", encoding="utf-8") as f: - code = f.read() - - ctx = build_exec_context(types=self._types, - predicates=self._get_all_predicates(), - options=self._get_all_options(), - extra_context={ - "params": - _ParamsView(self._fitted_params), - "ParamSpec": ParamSpec, - }) - - option_names = {o.name for o in self._get_all_options()} - valid, warnings, err = load_learned_samplers(code, ctx, option_names) - if err is not None: - logger.warning("Failed to load %s:\n%s", path, err) - return {} - for warning in warnings: - logger.warning("%s: %s", path, warning) - return valid - - def _synthesize_samplers_standalone( - self, trajectories: List[LowLevelTrajectory], - base_pred_triples: List[Tuple[State, Action, State]], - inferred_hint: Dict[str, List[str]]) -> None: - """Run a dedicated sampler-synthesis session. - - Used when oracle_sim_program short-circuits the sim-synthesis - session, so samplers still get learned. Reuses that session's - sandbox/snapshot/tool machinery. Called from _learn_simulator - after the option model is built, so the session's probe has a - working simulator. - """ - if CFG.agent_sdk_use_local_sandbox: - sandbox_dir: Optional[str] = os.path.abspath( - os.path.join(self._get_log_dir(), "sandbox")) - else: - sandbox_dir = self._tool_context.sandbox_dir - base = sandbox_dir or self._get_log_dir() - - if CFG.agent_sdk_use_local_sandbox: - sandbox_dir_for_agent: Optional[str] = "." - elif sandbox_dir: - sandbox_dir_for_agent = "/sandbox" - else: - sandbox_dir_for_agent = None - - paths = self._sampler_paths(base) - simulator_file = os.path.join(base, "simulator.py") - versions_dir = os.path.join(base, "simulator_versions") - - # Same namespace the main synthesis session gets (trajectories, - # train_tasks, is_goal_state, describe_trajectory, np, ParamSpec, - # evaluate_trajectory when the env defines evaluators). - exec_ns: Dict[str, Any] = self._build_synthesis_exec_ns(trajectories) - # The probe's `sim.refine` gives the agent the samples-to-refine - # feedback signal; the sampler tool installs + sanity-checks the - # samplers. - # pylint: disable-next=import-outside-toplevel - from predicators.agent_sdk.belief_probe import _check_time_budget - toolkit = create_synthesis_tools( - exec_ns, - base_pred_triples, - inferred_hint, - simulator_file=simulator_file, - versions_dir=versions_dir, - # The host class (AgentSimLearningApproach) provides the - # full backend surface; the mixin's own type covers only - # the sampler slice. - approach=cast("SynthesisBackend", self), - sandbox_dir=base, - sandbox_dir_for_agent=sandbox_dir_for_agent, - cycle_index_provider=self._learning_cycle_index, - budget_check=lambda: _check_time_budget(self._tool_context), - ) - tools = list(toolkit.tools) - self._install_sampler_surface(paths) - # Use the same declared surface as the mixin will assert against - # (_get_synthesis_tool_names already includes the sampler tool since - # _do_synthesize_samplers is True here). The rule-fitting surface is - # exposed but irrelevant — the message steers the agent to samplers. - declared = set(self._get_synthesis_tool_names() or ()) - self._tool_context.extra_mcp_tools = [ - t for t in tools if getattr(t, "name", "") in declared - ] - # The probe here runs the DEPLOYED belief model (no candidate - # provider: ctx.option_model already wraps the oracle sim - # program), which is exactly what samplers must speed up. The - # fit runner still targets simulator.py for consistency. - self._tool_context.probe_fit_provider = toolkit.fit_runner - self._tool_context.probe_validation_provider = toolkit.validation_runner - self._tool_context.probe_residuals_provider = \ - toolkit.residuals_runner - # pylint: disable-next=import-outside-toplevel - from predicators.agent_sdk.belief_probe import build_probe_namespace - probe_ns = build_probe_namespace(self._tool_context) - exec_ns["sim"] = probe_ns["sim"] - exec_ns["BeliefProbe"] = probe_ns["BeliefProbe"] - self._learning_mode = True - self._tool_context.extra_session_hooks = ( - self._build_synthesis_session_hooks( - [self._sampler_snapshot_target(paths)], base)) - - self._close_agent_session() - self._ensure_agent_session() - - predicate_listing = self._format_predicate_signatures( - self._get_all_predicates()) - options_digest = render_options_digest( - self._tool_context.options, - gt_options_ref_path=self._tool_context.gt_options_ref_path) - message = f"""\ -Synthesize per-skill samplers for this environment's options. The \ -simulator dynamics are already fixed (oracle/learned); your only job is \ -to make backtracking refinement land each option's continuous parameters \ -on its sketch-step subgoal instead of drawing them uniformly. - -## Available Predicates (subgoal geometry) -{predicate_listing} - -## Options -{options_digest} - -Explore the trajectory data with `run_python` (variables: \ -`trajectories`, `train_tasks`, `is_goal_state`, \ -`describe_trajectory(traj_idx)`, `np`, `ParamSpec`, plus the `sim` \ -probe over the deployed simulator - `sim.refine` is your \ -samples-to-refine feedback signal).""" - message = message + "\n\n" + self._sampler_synthesis_message(paths) - - try: - responses = self._query_agent_sync(message, kind="learn") - dead = query_fatal_error(responses) - if dead is not None: - # Nothing was synthesized: stop before this cycle is - # checkpointed as learned (see the simulator learn). - raise AgentSessionFatalError( - "The sampler-synthesis session died without the agent " - f"doing any work ({dead}); refusing to checkpoint this " - "cycle as learned.") - finally: - self._tool_context.extra_session_hooks = {} - self._tool_context.extra_mcp_tools = [] - self._tool_context.probe_artifact_loaders.clear() - self._tool_context.probe_fit_provider = None - self._tool_context.probe_validation_provider = None - self._tool_context.probe_residuals_provider = None - self._learning_mode = False - self._close_agent_session() - - self._finalize_and_load_samplers(paths) diff --git a/predicators/approaches/synthesis_validation.py b/predicators/approaches/synthesis_validation.py index 4a759fc6ec..e518fed271 100644 --- a/predicators/approaches/synthesis_validation.py +++ b/predicators/approaches/synthesis_validation.py @@ -16,10 +16,8 @@ from typing import TYPE_CHECKING, Any, Dict, List, Tuple from predicators.code_sim_learning.fit_space import ParamSpec -from predicators.code_sim_learning.fitting import fit_rule_parameters from predicators.code_sim_learning.utils import LearnedSimulator, \ - apply_rules, has_latent_rules, has_physics_rules -from predicators.structs import Action, State + apply_rules, has_latent_rules if TYPE_CHECKING: from predicators.agent_sdk.synthesis_backend import SynthesisBackend @@ -31,26 +29,20 @@ def build_candidate_option_model( approach: "SynthesisBackend", rules: List, specs: List[ParamSpec], - residual_features: Dict[str, List[str]], - base_pred_triples: List[Tuple[State, Action, State]], latent_init: Any = None, - fit: bool = True, -) -> Tuple[Any, Dict[str, float], float]: - """Fit ``specs`` (unless ``fit=False``) and build the candidate's option - model. - - ``fit=False`` builds the candidate at :func:`carry_over_params` - (the last published fit where a spec still exists and the value - lies in its box, the declared init value otherwise) and returns - ``nan`` for the SSE: fitting is the agent's explicit ``sim.fit`` - call, never a side effect of probing (see +) -> Tuple[Any, Dict[str, float]]: + """Build the candidate's option model at :func:`carry_over_params`. + + The parameters are the last published fit's where a spec still + exists and the value lies in its box, the declared init value + otherwise: fitting is the agent's explicit ``sim.fit`` call, never a + side effect of probing (see ``AgentSimLearningApproach._make_candidate_probe_model_provider``). The front half of the synthesis-session probe: every rollout must exercise the candidate simulator at its *deployed* (fitted) parameters, never at init_value. Returns ``(option_model, - fitted_params, fit_sse)``; raises ``RuntimeError`` when fitting - fails. + params)``. Publishes side effects onto ``approach`` exactly once, here, so the two surfaces can never disagree: the candidate ``rules`` / @@ -59,12 +51,6 @@ def build_candidate_option_model( place* (invented predicates hold a ``_ParamsView`` over it - the gating rule and the gating predicate must anchor to the same values). - - Recurrent (latent-declaring, 5-arg) rules are fit with the latent - threaded per trajectory; fully-observable rules take the legacy - per-transition path. Dispatch keys off the candidate rule - signatures (:func:`has_latent_rules`), as everywhere else in the - fitting stack. """ # pylint: disable=protected-access latent = has_latent_rules(rules) @@ -80,37 +66,10 @@ def build_candidate_option_model( if latent: approach._latent_init = latent_init - if not fit: - params = carry_over_params(approach._fitted_params, specs) - approach._fitted_params.clear() - approach._fitted_params.update(params) - return _finish_candidate_model(approach, rules, params), params, \ - float("nan") - try: - if has_physics_rules(rules): - # Physics-command rules act through engine stepping, so the - # teacher-forced objectives below cannot see them; fit - # against free-running rollouts instead (the same routing - # sim.fit uses). The joint fit also covers any declared - # PHYSICAL_PARAM_SPECS, which _load_simulator_from_module_file - # published onto the approach before this runs. - fit_result, fit_sse = approach._fit_parameters_joint_rollout( - rules, specs, residual_features) - elif latent: - fit_result, fit_sse = approach._fit_parameters_recurrent( - rules, specs, base_pred_triples, residual_features) - else: - fit_result, fit_sse = fit_rule_parameters(rules, specs, - base_pred_triples, - residual_features) - params = fit_result.point_estimate - except Exception as e: - raise RuntimeError(f"param fitting failed:\n{e}") from e - - # In place (clear + update, never replace): see docstring. + params = carry_over_params(approach._fitted_params, specs) approach._fitted_params.clear() approach._fitted_params.update(params) - return _finish_candidate_model(approach, rules, params), params, fit_sse + return _finish_candidate_model(approach, rules, params), params def carry_over_params(fitted: Dict[str, float], diff --git a/predicators/envs/pybullet_domino/env.py b/predicators/envs/pybullet_domino/env.py index e95ab76d35..ce62bcd36b 100644 --- a/predicators/envs/pybullet_domino/env.py +++ b/predicators/envs/pybullet_domino/env.py @@ -270,23 +270,18 @@ def _configure_instance_physics(self) -> None: # automatically after every reset_state. friction = CFG.domino_true_friction if self._skip_domain_specific_dynamics and \ - CFG.domino_planning_friction is not None and \ - not CFG.agent_sim_learn_oracle_sim_params: - # agent_sim_learn_oracle_sim_params grants the planner the - # TRUE friction (oracle upper bound) while task generation keeps - # using domino_planning_friction for the differentiation filter. + CFG.domino_planning_friction is not None: friction = CFG.domino_planning_friction if self._domino_component is not None and abs( friction - self._domino_component.domino_friction) > 1e-9: self.set_domino_physical_params(lateral_friction=friction) # Heavy-block tasks: planning sims BELIEVE the heavy gray blocks # are ordinary dominoes (normal mass), so their rollouts propagate - # a chain straight through one. The eval env (and the oracle- - # params planner) keeps the true heavy mass, asserted at reset. + # a chain straight through one. The eval env keeps the true heavy + # mass, asserted at reset. if CFG.domino_heavy_block_tasks \ and self._domino_component is not None \ - and self._skip_domain_specific_dynamics \ - and not CFG.agent_sim_learn_oracle_sim_params: + and self._skip_domain_specific_dynamics: self.set_domino_physical_params(block_mass=self.domino_mass) def _create_robot_predicates(self) -> None: diff --git a/predicators/execution_monitoring/subgoal_annotations_monitor.py b/predicators/execution_monitoring/subgoal_annotations_monitor.py deleted file mode 100644 index 93659e817e..0000000000 --- a/predicators/execution_monitoring/subgoal_annotations_monitor.py +++ /dev/null @@ -1,99 +0,0 @@ -"""An execution monitor that checks plan-sketch subgoal annotations at option -boundaries and suggests replanning on divergence.""" - -import logging -from dataclasses import dataclass -from typing import Any, Optional, Sequence - -from predicators.execution_monitoring.base_execution_monitor import \ - BaseExecutionMonitor -from predicators.structs import State, _Option - - -@dataclass -class SubgoalExecutionStatus: - """Live execution status of an annotated plan, exported by an approach via - ``get_execution_monitoring_info``. - - ``sketch`` items are duck-typed sketch steps exposing - ``subgoal_atoms``, ``subgoal_neg_atoms`` and ``option`` (see - ``agent_sdk.bilevel_sketch.SketchStep``); the type is kept loose so - the monitoring layer does not import agent_sdk. The approach's - dispensed policy mutates ``steps_initiated``/``current_option`` as - it executes, so the monitor always sees the live values. - """ - sketch: Sequence[Any] - steps_initiated: int = 0 - current_option: Optional[_Option] = None - - -class SubgoalAnnotationsExecutionMonitor(BaseExecutionMonitor): - """Suggest replanning when the step that just finished has a subgoal - annotation that does not hold in the real state. - - The check happens at the exact option boundary: when the currently - executing option's terminal condition is true in the given state, - the step it completes is checked before the policy advances to the - next option. Forward validation only proves a plan works in the - option model; real execution can still diverge (e.g. a place whose - drop-settle is chaotic lands off-target), after which the remaining - open-loop plan is doomed — it burns the episode horizon waiting for - effects that can no longer occur. Two boundaries are not caught: - divergence that only manifests inside a non-terminating option, and - Wait steps ended by the atom-change path in - ``utils.option_policy_to_policy`` (those terminate exactly when - their target atoms — derived from the same annotation — hold, so the - check would pass anyway). - """ - - @classmethod - def get_name(cls) -> str: - return "subgoal_annotations" - - def step(self, state: State) -> bool: - # No active annotated plan (e.g. exploration with an override - # policy, or replanning disabled): never suggest replanning. - if not self._approach_info: - return False - status = self._approach_info[0] - if not isinstance(status, SubgoalExecutionStatus): - return False - option = status.current_option - if option is None or status.steps_initiated <= 0: - return False - # Note: terminal() is also called by the policy machinery on - # this same state; skill terminal functions read memory but do - # not mutate it, so the double call is safe. - if not option.terminal(state): - return False - step_idx = status.steps_initiated - 1 - step = status.sketch[step_idx] - - def _holds(atom: Any) -> Optional[bool]: - # An atom the classifier cannot evaluate on this observation - # (e.g. an invented classifier that indexes a latent real - # states do not carry) is unverifiable, not diverged: skip - # it with a warning instead of crashing the episode or - # aborting on a claim no observation can settle. - try: - return bool(atom.holds(state)) - except Exception as e: # pylint: disable=broad-except - logging.warning( - "Subgoal atom %s is unverifiable on the real " - "observation (%s: %s); skipping it.", atom, - type(e).__name__, e) - return None - - unsat = [ - str(a) for a in (step.subgoal_atoms or set()) if _holds(a) is False - ] - unsat += [ - f"NOT {a}" for a in (step.subgoal_neg_atoms or set()) - if _holds(a) is True - ] - if not unsat: - return False - logging.info( - "Subgoal divergence after step %d (%s): unsatisfied %s. " - "Suggesting replan.", step_idx, step.option.name, sorted(unsat)) - return True diff --git a/predicators/explorers/__init__.py b/predicators/explorers/__init__.py index 31995dffbc..3a930dbbe7 100644 --- a/predicators/explorers/__init__.py +++ b/predicators/explorers/__init__.py @@ -1,13 +1,11 @@ """Handle creation of explorers.""" -import logging -from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Set +from typing import Callable, Dict, List, Optional, Set from gym.spaces import Box from predicators import utils from predicators.competence_models import SkillCompetenceModel -from predicators.explorers.agent_explorer_base import AgentExplorerBase from predicators.explorers.base_explorer import BaseExplorer from predicators.explorers.bilevel_planning_explorer import \ BilevelPlanningExplorer @@ -18,10 +16,6 @@ NSRTSamplerWithEpsilonIndicator, ParameterizedOption, Predicate, State, \ Task, Type, _GroundSTRIPSOperator -if TYPE_CHECKING: - from predicators.agent_sdk.session_manager import SessionManagerProtocol - from predicators.agent_sdk.tools import ToolContext - __all__ = ["BaseExplorer"] # Find the subclasses. @@ -49,23 +43,10 @@ def create_explorer( seen_train_task_idxs: Optional[Set[int]] = None, pursue_task_goal_first: Optional[bool] = None, maple_q_function: Optional[MapleQFunction] = None, - tool_context: Optional["ToolContext"] = None, - agent_session: Optional["SessionManagerProtocol"] = None, ) -> BaseExplorer: """Create an explorer given its name.""" if max_steps_before_termination is None: max_steps_before_termination = CFG.max_num_steps_interaction_request - # Deprecated aliases from before the explorers' model-free / - # model-based rename (2026-08-30); old launch commands and - # requeued jobs still pass them. - aliases = { - "agent_plan": "agent_model_free", - "agent_bilevel": "agent_model_based", - } - if name in aliases: - logging.warning("Explorer name %r is deprecated; use %r.", name, - aliases[name]) - name = aliases[name] for cls in utils.get_all_subclasses(BaseExplorer): if not cls.__abstractmethods__ and cls.get_name() == name: # Special case GLIB because it uses babble predicates and an atom @@ -122,13 +103,6 @@ def create_explorer( action_space, train_tasks, max_steps_before_termination, nsrts, maple_q_function) - elif issubclass(cls, AgentExplorerBase): - assert tool_context is not None - assert agent_session is not None - explorer = cls(initial_predicates, initial_options, types, - action_space, train_tasks, - max_steps_before_termination, tool_context, - agent_session) else: explorer = cls(initial_predicates, initial_options, types, action_space, train_tasks, diff --git a/predicators/explorers/agent_explorer_base.py b/predicators/explorers/agent_explorer_base.py deleted file mode 100644 index 90b7eee828..0000000000 --- a/predicators/explorers/agent_explorer_base.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Shared plumbing for the agent explorers. - -Both agent explorers query a Claude agent session for a plan and roll it -out in the real environment; they differ in what the agent is given (the -model-based explorer exposes the learned belief simulator through tools, -the model-free one only the task description). This base class holds the -common session wiring: construction, the trajectory summary shown to the -agent, final-text extraction, and the random-options fallback. -""" - -from typing import Any, Dict, List, Optional, Set - -from gym.spaces import Box - -from predicators import utils -from predicators.agent_sdk.learn_prompts import render_world_model_notes_block -from predicators.agent_sdk.response_parser import extract_final_text -from predicators.agent_sdk.session_manager import SessionManagerProtocol -from predicators.agent_sdk.sketch_prompts import summarize_trajectories -from predicators.agent_sdk.tools import ToolContext -from predicators.explorers.base_explorer import BaseExplorer -from predicators.structs import Action, ExplorationStrategy, \ - ParameterizedOption, Predicate, State, Task, Type - - -class AgentExplorerBase(BaseExplorer): - """Base class for explorers that query a Claude agent session.""" - - def __init__(self, predicates: Set[Predicate], - options: Set[ParameterizedOption], types: Set[Type], - action_space: Box, train_tasks: List[Task], - max_steps_before_termination: int, tool_context: ToolContext, - agent_session: SessionManagerProtocol) -> None: - super().__init__(predicates, options, types, action_space, train_tasks, - max_steps_before_termination) - self._tool_context = tool_context - self._agent_session = agent_session - - def _agent_tool_names(self) -> Optional[List[str]]: - """Return tool names exposed by the current session, if any.""" - return getattr(self._agent_session, "tool_names", None) - - def _random_options_fallback(self) -> ExplorationStrategy: - """Fall back to random option sampling.""" - - def fallback_policy(state: State) -> Action: - del state - raise utils.RequestActPolicyFailure( - "Random option sampling failed!") - - policy = utils.create_random_option_policy(self._options, self._rng, - fallback_policy) - return policy, lambda _: False - - def _build_trajectory_summary(self) -> str: - """Summarize trajectory data for the agent.""" - all_trajs = (self._tool_context.offline_trajectories + - self._tool_context.online_trajectories) - return summarize_trajectories(all_trajs, - self._predicates, - train_tasks=self._train_tasks) - - def _world_model_notes_block(self) -> str: - """The natural-language world model quoted into the explore prompt - (empty unless the approach keeps one).""" - return render_world_model_notes_block( - self._tool_context.world_model_notes, - self._tool_context.world_model_notes_path) - - @staticmethod - def _extract_option_plan_text(responses: List[Dict[str, Any]]) -> str: - """Extract plan text from the last assistant text response.""" - return extract_final_text(responses) diff --git a/predicators/explorers/agent_model_based_explorer.py b/predicators/explorers/agent_model_based_explorer.py deleted file mode 100644 index 2bb31d7285..0000000000 --- a/predicators/explorers/agent_model_based_explorer.py +++ /dev/null @@ -1,512 +0,0 @@ -"""Agent model-based explorer: the agent sketches an experiment against the -learned belief model, and it runs as written. - -Queries a Claude agent for a fully parameterized plan sketch and rolls -it out for real exactly as written. The agent refines and validates -in-session against the currently-learned belief model (``sim.refine``, -``sim.run``, ``submit_plan``); a plan that passed the capture -gate executes as a belief-certified solve attempt, anything else -executes as an experiment with no harness-side parameter search or -substitution. When the belief model disagrees with reality (e.g. a -subgoal atom it expected after a Wait doesn't actually hold), the -trajectory is a targeted learning signal for online simulator -synthesis. - -Parallels ``AgentModelBasedApproach`` for the sketch workflow; the -session plumbing lives in ``AgentExplorerBase``. - -Registered under the CLI explorer name ``agent_model_based`` -(``agent_bilevel`` is kept as a deprecated alias). -""" - -import logging -import os -from typing import Any, Callable, Dict, List, Optional, Sequence - -import numpy as np - -from predicators import utils -from predicators.agent_sdk import bilevel_sketch -from predicators.agent_sdk.rendering import save_task_state_image -from predicators.agent_sdk.session_base import AgentSessionFatalError, \ - query_fatal_error -from predicators.agent_sdk.session_manager import run_query_sync -from predicators.agent_sdk.tools import PlanCapture, agent_render_resolution, \ - load_ground_sampler_fns -from predicators.explorers.agent_explorer_base import AgentExplorerBase -from predicators.settings import CFG -from predicators.structs import Action, ExplorationStrategy, State, Task - - -class AgentModelBasedExplorer(AgentExplorerBase): - """Queries a Claude agent for a plan sketch and executes it as written.""" - - @classmethod - def get_name(cls) -> str: - return "agent_model_based" - - # ------------------------------------------------------------------ # - # Exploration strategy - # ------------------------------------------------------------------ # - - def _get_exploration_strategy(self, train_task_idx: int, - timeout: int) -> ExplorationStrategy: - task = self._train_tasks[train_task_idx] - # The approach syncs tool_context.option_model right before building - # this explorer, so reading here picks up the latest learned model. - option_model = self._tool_context.option_model - assert option_model is not None, \ - "agent_model_based explorer needs a synced option_model" - - # Reset the per-request mental-model verdict so a stale value can't - # leak if the query below throws or falls back to random before - # producing one. - self._tool_context.last_mental_model_solved = None - - # Point the agent's interactive tools (submit_plan, the - # sim probe) at the EXPLORE task. They - # default to ctx.current_task when the agent omits task_idx, and - # test-time _solve leaves current_task on 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 - # lacks), 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 submit_plan but 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 execute it at its captured params. 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 - # Exploration delivers plan sketches even in policy-mode configs: - # make sure a solve attempt's policy_capture_mode never bleeds - # into this query (it would disable the plan capture gate). - self._tool_context.policy_capture_mode = False - self._tool_context.clear_plan_capture() - - try: - prompt = bilevel_sketch.build_solve_prompt( - task, - all_predicates=self._predicates, - all_options=self._options, - trajectory_summary=self._build_trajectory_summary(), - tool_names=self._agent_tool_names(), - experiment_guidance=self._build_experiment_guidance(), - # Plans generated by this cycle's earlier requests: ask for a - # complementary plan instead of the identical one repeated. - scheduled_plans=list(self._tool_context.cycle_scheduled_plans), - initial_image_section=self._initial_image_section( - task, train_task_idx), - propose_params=CFG.agent_bilevel_use_llm_initial_params, - # The sketch is the experiment; a capture-gate-validated - # plan is welcome (it counts as a solve) but not required. - require_tool_validation=False, - # Explore contract: the sketch is a real-env experiment, and - # the belief model may lack goal-critical dynamics, so a - # simulator-failing sketch is a valid deliverable. - explore_mode=True, - ) - responses = run_query_sync(self._agent_session, - prompt, - kind="explore") - dead = query_fatal_error(responses) - if dead is not None: - # The session backend refused the query (usage limit, - # auth, transport): a random-options episode here would - # be junk data that the cycle then learns from (2026-08-28 - # run_20260827_171610 cycle 3). Terminate the run instead; - # the relaunch re-explores this cycle. - raise AgentSessionFatalError( - "explore query died without the agent doing any work " - f"({dead}); not falling back to random exploration.") - plan_text = self._extract_option_plan_text(responses) - # The session's tool capture: a goal-reaching plan the agent - # validated in the belief through the capture gate - # (submit_plan, N fresh - # rollouts). ``reached_goal`` is the gate's verdict. - capture = self._tool_context.take_plan_capture() - if CFG.agent_explorer_execute_certified_plan and capture.plan \ - and capture.reached_goal is True: - # Certified: the mental model solves the task with THIS - # plan, so run it verbatim as a solve attempt instead of - # re-searching (or boundary-probing) its parameters. A - # real success now counts for early stopping. The cycle's - # later requests see it under "plans already scheduled" - # and are asked for a DIFFERENT certified plan (a second - # test of the model), resubmitting this one only as a - # last resort. - plan = list(capture.plan) - logging.info( - "agent_model_based explorer: the agent's tool-validated " - "plan passed the belief's capture gate (%s); executing " - "it verbatim as this episode's solve attempt (mental " - "model solved the goal).", capture.validation_summary - or "goal reached") - if capture.sketch: - self._tool_context.last_sketch_subgoals = [ - (s.subgoal_atoms, s.subgoal_neg_atoms) - for s in capture.sketch - ] - self._tool_context.last_sketch_options = [ - (s.option.name, [o.name for o in s.objects]) - for s in capture.sketch - ] - self._tool_context.last_mental_model_solved = True - self._tool_context.cycle_scheduled_plans.append( - self._format_plan(plan) + - "\n NOTE: belief-certified; executes verbatim as a " - "solve attempt.") - return self._certified_plan_strategy(plan) - if not plan_text and not capture.plan: - raise ValueError("agent returned empty plan text") - - gs_fns, gs_err = load_ground_sampler_fns(self._tool_context) - if gs_err is not None: - logging.warning("[explore] %s", gs_err) - sketch = bilevel_sketch.parse_sketch_from_text( - plan_text, - task, - predicates=self._predicates, - options=self._options, - types=self._types, - parse_continuous_params=CFG. - agent_bilevel_use_llm_initial_params, - parse_ground_samplers=CFG.agent_bilevel_ground_samplers, - ground_sampler_fns=gs_fns or None, - ) if plan_text else [] - if not sketch: - sketch = self._sketch_from_capture(capture) or [] - if not sketch: - raise ValueError("parsed empty plan sketch") - self._tool_context.last_sketch_subgoals = [ - (s.subgoal_atoms, s.subgoal_neg_atoms) for s in sketch - ] - self._tool_context.last_sketch_options = [ - (s.option.name, [o.name for o in s.objects]) for s in sketch - ] - # The agent's sketch IS the experiment: it runs in the real - # environment exactly as written. The harness does no belief - # refinement of it - the agent refines and validates - # in-session (sim.refine / sim.run / submit_plan), - # and only a capture-gate-certified plan (handled above) - # counts as a belief-validated solve for early stopping. - plan = self._ground_sketch_verbatim(sketch) - self._tool_context.last_mental_model_solved = False - record = self._format_sketch(sketch, plan) - logging.info( - "agent_model_based explorer: executing the agent's sketch " - "verbatim for train task %d (%d steps; not " - "belief-certified):\n%s", train_task_idx, len(plan), record) - self._tool_context.cycle_scheduled_plans.append( - record + "\n NOTE: executes as written, without " - "belief-model certification; a real success will NOT " - "count toward early stopping.") - policy = utils.option_plan_to_policy( - plan, - abstract_function=lambda s: utils.abstract( - s, self._predicates)) - return self._wrap_policy(policy), lambda _: False - except AgentSessionFatalError: - # A random fallback would hide the broken session backend; - # re-raise so the run terminates. - raise - except Exception as e: # pylint: disable=broad-except - logging.warning(f"agent_model_based explorer failed: {e}. " - "Falling back to random options.") - - if not CFG.agent_explorer_fallback_to_random: - raise utils.RequestActPolicyFailure( - "agent_model_based explorer failed and fallback disabled.") - return self._random_options_fallback() - - # ------------------------------------------------------------------ # - # Helpers - # ------------------------------------------------------------------ # - - def _sketch_from_capture( - self, - capture: PlanCapture) -> Optional[List[bilevel_sketch.SketchStep]]: - """Rebuild a sketch from a captured, tool-validated plan, or None. - - ``submit_plan`` stashes a forward-validated 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 plan executes at exactly - the values the agent validated. The capture was already taken - (consumed) by the caller. - """ - plan = capture.plan - captured_sketch = capture.sketch - 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_model_based explorer: final text didn't parse, recovered " - "the " - "agent's tool-validated plan from capture (%d steps); executing " - "it at the captured params.", len(seeded)) - return seeded - - def _ground_sketch_verbatim( - self, sketch: Sequence[bilevel_sketch.SketchStep]) -> List[Any]: - """Ground each sketch step at the agent's proposed parameters. - - Nothing is searched or substituted: the executed plan is the - agent's. A step left without parameters (or with the wrong - arity) gets ONE uniform draw from the option's box and a - warning. Wait steps carry their annotated subgoals as - ``wait_target_atoms`` so the option terminates on the intended - atom change. - """ - plan: List[Any] = [] - for i, step in enumerate(sketch): - dim = step.option.params_space.shape[0] - params = step.initial_params - if params is None or len(params) != dim: - if dim > 0: - logging.warning( - "agent_model_based explorer: step %d (%s) has no " - "usable proposed params (%s); drawing one sample " - "from the option's box - propose every " - "parameter explicitly.", i, step.option.name, - None if params is None else list(params)) - params = bilevel_sketch.sample_params(step.option, self._rng) - plan.append( - bilevel_sketch.ground_step( - step, np.asarray(params, dtype=np.float32))) - return plan - - @staticmethod - def _format_sketch(sketch: Sequence[bilevel_sketch.SketchStep], - plan: Sequence[Any]) -> str: - """One indented ``i: Option(objs)[params] -> {atoms}`` line per - grounded step (params as executed, atoms as annotated).""" - lines = [] - for i, (step, opt) in enumerate(zip(sketch, plan)): - objs = ", ".join(o.name for o in opt.objects) - par = ", ".join(f"{p:.4f}" for p in opt.params) - line = f" {i}: {opt.name}({objs})[{par}]" - atoms = sorted(str(a) for a in (step.subgoal_atoms or set())) - atoms += sorted(f"NOT {a}" - for a in (step.subgoal_neg_atoms or set())) - if atoms: - line += f" -> {{{', '.join(atoms)}}}" - lines.append(line) - return "\n".join(lines) - - @staticmethod - def _format_plan(plan: Sequence[Any]) -> str: - """One indented ``i: Option(objs)[params]`` line per grounded - option.""" - lines = [] - for i, opt in enumerate(plan): - obj_s = ", ".join(o.name for o in opt.objects) - par_s = ", ".join(f"{p:.4f}" for p in opt.params) - lines.append(f" {i}: {opt.name}({obj_s})[{par_s}]") - return "\n".join(lines) - - def _certified_plan_strategy(self, - plan: Sequence[Any]) -> ExplorationStrategy: - """Execute a belief-certified grounded plan verbatim.""" - logging.info("agent_model_based explorer: certified plan:\n%s", - self._format_plan(plan)) - policy = utils.option_plan_to_policy( - list(plan), - abstract_function=lambda s: utils.abstract(s, self._predicates)) - return self._wrap_policy(policy), lambda _: False - - def _wrap_policy( - self, policy: Callable[[State], - Action]) -> Callable[[State], Action]: - """Convert OptionExecutionFailure into RequestActPolicyFailure. - - Lets the main loop cleanly terminate the episode when the - refined plan finishes or fails mid-execution (which is exactly - the disagreement signal we want to collect). - """ - - def _wrapped(state: State) -> Action: - try: - return policy(state) - except utils.OptionExecutionFailure as e: - raise utils.RequestActPolicyFailure(e.args[0], e.info) from e - - return _wrapped - - def _initial_image_section(self, task: Task, train_task_idx: int) -> str: - """Render the explore task's initial state and return a prompt section - pointing at it, mirroring what test-time solves get. - - Saved as ``train_task{N:03d}_initial_state.png`` so train-task - scenes are inspectable alongside the test-task init images. - Empty string when rendering is unavailable (e.g. the sandbox - isn't created yet, so ``image_save_dir`` is unset). - """ - env = self._tool_context.env - save_dir = self._tool_context.image_save_dir - if env is None or save_dir is None: - return "" - img_name = f"train_task{train_task_idx:03d}_initial_state.png" - with agent_render_resolution(): - saved = save_task_state_image(env, task, save_dir, img_name) - if saved is None: - return "" - # cwd of the agent is the sandbox root, so reference test_images/. - return ("\n## Initial State Image\n" - "A rendering of the initial scene has been saved to " - f"`./test_images/{img_name}`. **Read this image first** to " - "understand the spatial layout before planning.\n") - - def _build_experiment_guidance(self) -> str: - """LLM-proposal half of active-experiment design. - - Always injects the learn phase's open-questions ledger (the - ranked experiment specs it wrote for exploration to run) when - one exists in the sandbox. When info-seeking is on, additionally - point the agent at ``sim.suggest_probes`` and - when an ensemble - scorer is wired - at the predicates the learned model is - currently most internally uncertain about. - """ - parts = [] - ledger = self._read_open_questions() - if ledger: - parts.append( - "The learning phase left this ranked ledger of OPEN " - "QUESTIONS - uncertainties it could not settle from the " - "data collected so far, each with the experiment that " - "would settle it. The TOP entry is mandatory for this " - "cycle: run its experiment as specified (its option " - "sequence and parameters) unless a plan already " - "scheduled this cycle covers it, and fold in as many " - "lower entries as the step budget allows:\n" + ledger) - if self._tool_context.info_seeking_active(): - parts.append( - "Your explicit continuous parameters execute exactly as " - "written. To find the parameters a step could be run at " - "to teach the model most, call " - "`sim.suggest_probes(plan_text)`: it rolls your sketch " - "forward on your own parameters and, per annotated step, " - "ranks feasible alternatives by the learned model's " - "ensemble disagreement on the step's subgoal atoms. Adopt " - "one by writing it into your sketch, only on a step whose " - "failure the episode can afford; annotate steps with the " - "geometry/timing you are least sure the model has right.") - disagreement = self._build_disagreement_summary() - if disagreement: - parts.append(disagreement) - # System-ID gaps from the previous learn phase (synced by the - # sim-learning approach): what the collected data could NOT - # support, phrased as experiment objectives. Exploration is the - # only place those gaps can be filled. - sysid = getattr(self._tool_context, "sysid_diagnostics", None) - if sysid: - parts.append( - "The previous system-identification fit left gaps that " - "only new interaction data can close:\n" + sysid) - return "\n\n".join(parts) - - _MAX_OPEN_QUESTIONS_CHARS = 4000 - - def _read_open_questions(self) -> str: - """The learn phase's ./open_questions.md ledger, or "". - - The ledger is ranked, so on overflow the head is kept. Read - directly from the sandbox instead of asking the session to go - find it: the explore query should START from the ledger, not - spend its budget rediscovering it. - """ - sandbox_dir = self._tool_context.sandbox_dir - if not sandbox_dir: - return "" - path = os.path.join(sandbox_dir, "open_questions.md") - try: - with open(path, "r", encoding="utf-8") as f: - text = f.read().strip() - except OSError: - return "" - if len(text) > self._MAX_OPEN_QUESTIONS_CHARS: - # Cut at an entry boundary, never mid-sentence, and point at - # the file: a mid-sentence cut silently dropped the entries - # the header told the agent to fold in (run_20260830). - head = text[:self._MAX_OPEN_QUESTIONS_CHARS] - cut = head.rfind("\n#") - if cut <= 0: - cut = head.rfind("\n\n") - if cut > 0: - head = head[:cut] - text = (head.rstrip() + - "\n[... ledger truncated at the prompt cap - read " - "./open_questions.md for the remaining entries]") - return text - - def _build_disagreement_summary(self) -> str: - """Name the predicates the ensemble disagrees most about. - - Scans a bounded sample of recent-trajectory states, scores each - abstract atom's ensemble disagreement via the wired scorer, and - reports the highest-disagreement predicates. Grounded in the - actual ensemble, so it points the agent at genuinely-uncertain - dynamics rather than guesses. Empty when no scorer/trajectories. - """ - fn = self._tool_context.atom_disagreement_fn - if fn is None: - return "" - all_trajs = (self._tool_context.offline_trajectories + - self._tool_context.online_trajectories) - if not all_trajs: - return "" - recent = all_trajs[-CFG.agent_sdk_max_trajectories_in_context:] - states: List[State] = [] - for traj in recent: - n = len(traj.states) - if n == 0: - continue - stride = max(1, n // 6) # <= ~6 states/trajectory to bound cost - states.extend(traj.states[::stride]) - best: Dict[str, float] = {} - for s in states: - for atom in utils.abstract(s, self._predicates): - try: - d = float(fn(s, {atom})) - except Exception: # pylint: disable=broad-except - continue - name = atom.predicate.name - if d > best.get(name, 0.0): - best[name] = d - # One log line with the full ranking (scope note: abstract() yields - # true atoms only, so a predicate absent here was never measured, not - # necessarily agreed-upon). All values <= 0.05 -> no guidance: the - # ensemble is internally confident (or too tight) everywhere. - all_ranked = sorted(((v, k) for k, v in best.items()), reverse=True) - logging.info( - "agent_model_based explorer: per-predicate max ensemble " - "disagreement " - "over %d states - %s.", len(states), - ", ".join(f"{k}={v:.4f}" for v, k in all_ranked) or "(none)") - ranked = [(v, k) for v, k in all_ranked if v > 0.05][:4] - if not ranked: - return "" - named = ", ".join(f"{k} (disagreement {v:.2f})" for v, k in ranked) - return ("Across recent trajectories, the learned model is most " - f"internally uncertain about: {named}. A sketch that puts " - "these predicates on the critical path will be most " - "informative.") diff --git a/predicators/explorers/agent_model_free_explorer.py b/predicators/explorers/agent_model_free_explorer.py deleted file mode 100644 index bfcfd04cbb..0000000000 --- a/predicators/explorers/agent_model_free_explorer.py +++ /dev/null @@ -1,175 +0,0 @@ -"""Agent model-free explorer: Claude agent generates grounded option plans -without a learned world model. - -Produces fully-grounded option plans (including continuous parameters) -from a one-shot task description and rolls them out in the real -environment. Unlike ``AgentModelBasedExplorer``, the agent has no belief -simulator to validate against and no backtracking refinement; it must -supply complete parameters itself. - -Registered under the CLI explorer name ``agent_model_free`` -(``agent_plan`` is kept as a deprecated alias). -""" - -import logging - -import numpy as np - -from predicators import utils -from predicators.agent_sdk.session_base import AgentSessionFatalError, \ - query_fatal_error -from predicators.agent_sdk.session_manager import run_query_sync -from predicators.explorers.agent_explorer_base import AgentExplorerBase -from predicators.settings import CFG -from predicators.structs import ExplorationStrategy, Task - - -class AgentModelFreeExplorer(AgentExplorerBase): - """Queries a Claude agent to produce grounded option plans.""" - - @classmethod - def get_name(cls) -> str: - return "agent_model_free" - - def _get_exploration_strategy(self, train_task_idx: int, - timeout: int) -> ExplorationStrategy: - task = self._train_tasks[train_task_idx] - try: - prompt = self._build_exploration_prompt(train_task_idx) - responses = run_query_sync(self._agent_session, - prompt, - kind="explore") - dead = query_fatal_error(responses) - if dead is not None: - # See agent_model_based_explorer: never explore at random - # because the session backend is down. - raise AgentSessionFatalError( - "explore query died without the agent doing any work " - f"({dead}); not falling back to random exploration.") - plan_text = self._extract_option_plan_text(responses) - if plan_text: - option_plan = self._parse_and_ground_plan(plan_text, task) - if option_plan: - # The Wait wrapper needs the abstraction to end a - # Wait on its target atoms. - policy = utils.option_plan_to_policy( - option_plan, - abstract_function=lambda s: utils.abstract( - s, self._predicates)) - return policy, lambda _: False - logging.info("Agent explorer: no valid plan, falling back to " - "random options.") - except AgentSessionFatalError: - # A random fallback would hide the broken session backend; - # re-raise so the run terminates. - raise - except Exception as e: # pylint: disable=broad-except - logging.warning(f"Agent explorer failed: {e}. " - "Falling back to random options.") - - if not CFG.agent_explorer_fallback_to_random: - raise utils.RequestActPolicyFailure( - "Agent explorer failed and fallback disabled.") - return self._random_options_fallback() - - def _build_exploration_prompt(self, train_task_idx: int) -> str: - """Build a prompt for the agent to produce an option plan.""" - task = self._train_tasks[train_task_idx] - init_state = task.init - - objects = list(init_state) - obj_strs = [] - for obj in sorted(objects, key=lambda o: o.name): - obj_strs.append(f" {obj.name}: {obj.type.name}") - - # Goal atoms - goal_strs = [str(a) for a in sorted(task.goal, key=str)] - - # Available options with signatures. - all_options = self._options - option_strs = [] - for opt in sorted(all_options, key=lambda o: o.name): - type_sig = ", ".join(t.name for t in opt.types) - params_dim = opt.params_space.shape[0] - if params_dim > 0: - low = opt.params_space.low.tolist() - high = opt.params_space.high.tolist() - param_info = (f", params_dim={params_dim}, " - f"low={low}, high={high}") - else: - param_info = "" - option_strs.append(f" {opt.name}({type_sig}{param_info})") - - # Current atoms - atoms = utils.abstract(init_state, self._predicates) - atom_strs = [str(a) for a in sorted(atoms, key=str)] - - # Trajectory summary - traj_summary = self._build_trajectory_summary() - - # Available tools - tools_str = "" - tool_names = self._agent_tool_names() - if tool_names: - tool_list = "\n".join(f" - {t}" for t in tool_names) - tools_str = f"\n## Available Tools\n{tool_list}\n" - - task_intro = ("You are exploring a task environment. " - f"Generate an option plan to explore task " - f"{train_task_idx}.") - prompt = f"""{task_intro} - -## Goal -{chr(10).join(goal_strs)} - -## Initial State Atoms -{chr(10).join(atom_strs)} - -## Objects -{chr(10).join(obj_strs)} - -## Available Options -{chr(10).join(option_strs)} -{traj_summary}{self._world_model_notes_block()}{tools_str} -## Instructions -Use your available tools to inspect the environment and test your plan before committing to it. - -Output an option plan, one option per line, in this exact format: -OptionName(obj1:type1, obj2:type2)[param1, param2] - -If an option has no continuous parameters, use empty brackets: OptionName(obj1:type1)[] - -Output ONLY the option plan lines at the end, after any analysis.""" - - return prompt - - def _parse_and_ground_plan(self, plan_text: str, task: Task) -> list: - """Parse option plan text and ground into executable options.""" - objects = list(task.init) - all_options = self._options - parsed = utils.parse_model_output_into_option_plan( - plan_text, - objects, - self._types, - all_options, - parse_continuous_params=True) - if not parsed: - logging.info("Agent explorer: parsed empty option plan.") - return [] - - grounded = [] - for option, objs, params in parsed: - try: - ground_opt = option.ground(objs, - np.array(params, dtype=np.float32)) - grounded.append(ground_opt) - except Exception as e: # pylint: disable=broad-except - logging.info(f"Agent explorer: failed to ground " - f"option {option.name}: {e}") - break - - if not grounded: - logging.info("Agent explorer: no options successfully grounded.") - else: - logging.info(f"Agent explorer: grounded {len(grounded)} options.") - return grounded diff --git a/predicators/settings.py b/predicators/settings.py index c1339787a4..116ab917be 100644 --- a/predicators/settings.py +++ b/predicators/settings.py @@ -1585,17 +1585,6 @@ class GlobalSettings: gnn_use_validation_set = True # parameters for GNN option policy approach - # GNN dynamics + shooting baseline (gnn_dynamics_shooting, paper arm - # C5): how many previous pre-option states ride along as node - # features (so a hidden mechanism is inferable from the recent - # past), the longest option sequence one shooting try samples, how - # many tries a plan query gets before failing, and whether the plan - # is re-shot from the observed state after every option (MPC) or - # executed open-loop. - gnn_dynamics_history_len = 2 - gnn_dynamics_max_plan_length = 30 - gnn_dynamics_shooting_max_tries = 200 - gnn_dynamics_replan_every_option = True gnn_option_policy_solve_with_shooting = True gnn_option_policy_shooting_variance = 0.1 gnn_option_policy_shooting_max_samples = 100 @@ -2117,17 +2106,11 @@ class GlobalSettings: agent_sdk_max_buffer_size = 20 * 1024 * 1024 agent_sdk_resume_session = True # resume previous session if available agent_sdk_max_trajectories_in_context = 3 - agent_sdk_log_agent_responses = True # Sandbox settings for agent SDK - agent_sdk_use_docker_sandbox = False # run agent inside Docker container - agent_sdk_docker_image = "predicators-sandbox" # Docker image name - # sandbox dir with built-in tools, no Docker + # sandbox dir with built-in tools agent_sdk_use_local_sandbox = False - # Agent explorer settings - agent_explorer_fallback_to_random = True # fall back to random on failure - # Agent planner approach settings agent_planner_use_scratchpad = False # include notes.md scratchpad # Whether the planner is given a simulator to test candidate plans with @@ -2164,40 +2147,6 @@ class GlobalSettings: # expressible). Default False hides the grammar from the agent and # rejects the annotations, keeping baseline arms free of the channel. agent_bilevel_ground_samplers = False - # When True, close the agent SDK session at the start of each test task - # so every test solve begins with a FRESH conversation (no context from - # earlier test tasks). The sandbox filesystem and learned artifacts are - # untouched. Default False keeps the current behavior: all test tasks - # share one continuous agent conversation. - agent_fresh_session_per_test_task = False - # Restart loop for test-task solving. Solve-time outcomes are close to - # heavy-tailed in agent-search quality (run_20260717 family split: the - # same tasks solved in 9-32 min in one launch and burned 2-11 h without - # solving in its identical sibling, anchored on wrong conclusions), so - # several short, independent attempts beat one long one. Each attempt - # above the first starts from a fresh conversation; the solve journal - # (below) carries curated knowledge across attempts. An attempt ends - # early with a validated (evaluator-solved) capture; otherwise its - # best-effort capture is banked and the best across attempts executes. - # Each attempt is exactly ONE agent query: however that query ends - - # a spent budget, an unparseable sketch, or a session that simply - # never submitted - the fresh-context restart is the only retry, so - # this is the sole knob controlling how many shots a task gets. Only - # the final attempt (no restart left) pays for the best-effort - # submission nudge. - agent_solve_max_attempts = 1 - # Wall-clock budget per solve attempt, in seconds (0 disables). The - # turn cap bounds turns, not compute - one run_python sweep hid - # 47k rollouts (~7 h) inside a single turn. On expiry, exploration - # tools refuse with a submit-now message and the approach runs the - # same best-effort submission flow as turn-cap exhaustion. - agent_solve_attempt_wall_clock = 0.0 - # When True, every solve attempt (including the first, i.e. every test - # task) begins with a fresh agent conversation; cross-attempt and - # cross-task knowledge travels through the solve journal instead of - # raw transcript history, which also carries the *wrong* conclusions - # of failed attempts. - agent_solve_fresh_context = False # Persistent per-run solve journal: the harness logs each attempt's # outcome + captured plan to /attempts.md, the agent keeps # its own lessons in /journal.md with the file tools, and @@ -2237,10 +2186,6 @@ class GlobalSettings: # cannot see (nothing fails). 3 tolerates a benign settle-in-place # step without letting a livelock burn the budget. agent_policy_max_repeated_noops = 3 - # LLM-free bypass: path to a prewritten policy.py used as the captured - # artifact for every test task (mirrors the sketch-file bypass). For - # smoke tests and debugging the execution path. - agent_policy_file = "" # --auto_resume only resumes from checkpoints modified within this # many hours. The checkpoint path ignores the run timestamp, so a # RELAUNCH of a finished experiment under the same experiment_id @@ -2287,22 +2232,8 @@ class GlobalSettings: # open-loop execution). Requires --execution_monitor # subgoal_annotations (enforced at approach construction). agent_bilevel_max_execution_replans = 0 - # When an execution replan's suffix refinement fails, whether to fall - # back to querying the agent for a fresh sketch - a brand-new - # full-turn-budget session. Default False: the cheap suffix replan is - # the only recovery, and when no suffix of the executed sketch - # refines from the diverged state the remaining plan resumes - # open-loop (the divergence is logged; the goal check decides the - # episode). Re-opening the agent budget is especially wasteful after - # a best-effort (non-solve) capture, whose execution diverges by - # construction. - agent_bilevel_replan_agent_fallback = False # log state pretty_str before/after each step agent_bilevel_log_state = False - # Load a plan sketch from a file instead of querying the LLM. The dir is - # under scripts/; the file may be a bare name or an absolute path. - agent_bilevel_plan_sketch_dir = "plan_sketches" - agent_bilevel_plan_sketch_file = "" # When a sketch refinement runs without an explicit timeout, the # caller computes # max(_min, _per_step * len(sketch)) @@ -2457,18 +2388,6 @@ class GlobalSettings: # Ensemble size used to estimate disagreement. 1 disables scoring # (every candidate scores 0) and reduces to first-feasible. agent_explorer_info_ensemble_size = 6 - # A plan the explore session validated through the capture gate - # (submit_plan: goal reached in - # agent_plan_validation_rollouts fresh belief rollouts) is executed - # verbatim as the episode's solve attempt with mental_model_solved= - # True. The cycle's remaining requests on that task still query the - # agent, which sees the certified plan among the plans already - # scheduled and is asked for a different certified plan (a second, - # independent test of the belief), resubmitting the same one only as - # a last resort; every certified attempt solving for real satisfies - # the train-driven early-stop rule. Off feeds the capture into the - # experiment search as seeds instead. - agent_explorer_execute_certified_plan = True # Per-parameter jitter as a fraction of the ParamSpec box width, for # the uniform-fallback ensemble only (see calibrated flag below). agent_explorer_info_perturb_frac = 0.15 @@ -2478,16 +2397,6 @@ class GlobalSettings: # fit runs). agent_explorer_info_calibrated_ensemble = True - # Code sim-learning parameter fitting settings. - # Persist the raw rollout-fit trajectories (states + actions per - # recorded episode) to /fit_data/ at every cycle-level - # fit. The fit data otherwise lives only in memory, which made the - # wrong fits of run_20260724_232411 (lateral_friction 1.0358 / - # 0.3236 vs true 0.5) impossible to replay offline: approximate - # re-execution from logged plans cannot reproduce mid-episode - # replans or the warm-env recording context, the very channel - # suspected of corrupting the fits. Cost: one small pickle per fit. - code_sim_learning_persist_fit_data = True # Truncate each rollout-fit trajectory once the scored features have # settled (physical_sysid.truncate_settled_tail): keep everything up # to the last observed motion plus a margin, drop the static tail. @@ -2789,12 +2698,6 @@ class GlobalSettings: # such and its anchor (env-registry baseline) is kept instead of # the fitted value. 0 disables the screen. code_sim_learning_rollout_sensitivity_factor = 2.0 - # Cross-cycle consistency check on the final per-cycle fit: a param - # whose MAP moved more than this many combined posterior sigmas - # since the previous cycle's fit is flagged (and its "identified" - # verdict downgraded) - mutually-incompatible confident fits are - # the signature of an overconfident probe. 0 disables. - code_sim_learning_rollout_cross_cycle_sigma = 3.0 # Pooled-evidence arbitration of a cross-cycle conflict: when the # new fit is flagged (see above) but explains the fit's surviving # segments with an SSE at least this factor smaller than the @@ -2822,12 +2725,6 @@ class GlobalSettings: # per fit. code_sim_learning_warm_start_with_lm = True - # Sim-learning oracle flags (for ablation / debugging). - # When True, load GT residual rules instead of running agent synthesis. - # Parameters init_values are perturbed so the fit still has work to do. - agent_sim_learn_oracle_sim_program = False - # Relative scale for perturbing oracle parameter init_values before the fit. - agent_sim_learn_oracle_sim_param_noise_scale = 0.2 # Ablations A6+A7 combined ("no uncertainty"): when False, nothing # consumes a posterior over the model parameters. The physics-margin sigma # points are never built (so the capture gate's physics margin and @@ -2857,42 +2754,6 @@ class GlobalSettings: agent_program_belief_particles = 6 agent_program_kernel_bandwidth = 0.2 agent_program_score_max_examples = 3 - # Ablation A2 ("zero-shot synthesis"): when True, the synthesis - # session runs even when no transition has been recorded, so the - # agent writes its artifacts from the task description, the scene - # and its own knowledge. Pair with no demos and - # num_online_learning_cycles 0 for one learn, one solve, done. - agent_sim_learn_zero_shot = False - # When True, use GT parameter values directly, skipping the fit. - # Also grants planning base sims the TRUE physical params (e.g. the true - # domino friction even when domino_planning_friction is set) — as if all - # param learning, rule-level and physical, had already succeeded. Task - # generation still reads domino_planning_friction for the - # differentiation filter, so the oracle, the no-learning baseline, and - # the sysID learner all see IDENTICAL tasks (and share the task cache: - # this agent_ flag is outside the cache key's - # domino_/pybullet_/skill_phase_ prefixes on purpose). - agent_sim_learn_oracle_sim_params = False - # When True, the agent learns PARAMETERIZED samplers - per-option - # (lifted-skill) functions that aim continuous option parameters at each - # sketch step's subgoal, instead of bilevel refinement drawing them - # uniformly from the option's box. The agent authors a versioned - # ``samplers.py`` (LEARNED_SAMPLERS keyed by option name) and tunes it - # with ``sim.samplers()``. Sampler learning rides along in - # the sim/predicate synthesis session when one runs - # (oracle_sim_program=False); when no synthesis session runs - # (oracle_sim_program=True) it gets a dedicated session of its own. - # The GROUND level of the sampler hierarchy needs no flag: a sketch - # step's ``~ [widths]`` region annotation compiles to a per-step - # GroundSampler that overrides the parameterized sampler for that step - # (ground > parameterized > uniform). - agent_sim_learn_parameterized_samplers = False - # When True (and parameterized_samplers is on), use ground-truth - # per-skill samplers from the env's GroundTruthSamplerFactory instead of - # having the agent learn them — if such samplers exist for the env; - # otherwise warn and fall back to synthesis. Mirrors - # agent_sim_learn_oracle_sim_program. - agent_sim_learn_oracle_samplers = False # Allowlist of env predicate names surfaced to the agent for # agent_sim_learning and its subclasses (e.g. diff --git a/predicators/structs.py b/predicators/structs.py index da9656d298..85b41650a3 100644 --- a/predicators/structs.py +++ b/predicators/structs.py @@ -2118,7 +2118,6 @@ class LowLevelTrajectory: _train_task_idx: Optional[int] = field(default=None) _source_simulator_version: Optional[str] = field(default=None) _source_predicates_version: Optional[str] = field(default=None) - _source_samplers_version: Optional[str] = field(default=None) _env_reward: Optional[float] = field(default=None) _env_terminated: Optional[bool] = field(default=None) @@ -2163,12 +2162,6 @@ def source_predicates_version(self) -> Optional[str]: collected this trajectory, or ``None`` if not tracked.""" return self._source_predicates_version - @property - def source_samplers_version(self) -> Optional[str]: - """Snapshot tag of the per-skill samplers used to generate the plan - that collected this trajectory, or ``None`` if not tracked.""" - return self._source_samplers_version - @property def env_rejected(self) -> bool: """Whether the supervisor (the environment's evaluator) rejected the @@ -2635,8 +2628,7 @@ class InteractionRequest: # explorers); online learning treats ``False`` as not-solved for # early stopping even if real-env execution happens to reach the # goal, so a model that executes-but-mispredicts isn't certified as - # trained. See AgentModelBasedExplorer / - # run.online_learning.generate_interaction_results. + # trained. See run.online_learning.generate_interaction_results. mental_model_solved: Optional[bool] = None diff --git a/predicators/utils.py b/predicators/utils.py index cf670e652b..03b9a65d0d 100644 --- a/predicators/utils.py +++ b/predicators/utils.py @@ -1612,18 +1612,13 @@ def __str__(self) -> str: def real_episode_step_budget(phase: Optional[str]) -> int: - """Low-level steps a real episode of this ``phase`` may use. - - Explore (interaction-request) episodes are capped by - ``max_num_steps_interaction_request`` on top of the horizon; test - and any other episodes by ``horizon`` alone. The belief tools quote - this number to the agent so its plans are sized for the budget the - real executor enforces (a 1000-step explore cap once went unstated - while the tools quoted the 3000-step horizon, and half the bridge - experiments were cut mid-plan). + """Low-level steps a real episode may use: the horizon, in every agent + session phase. + + The belief tools quote this number to the agent so its plans are + sized for the budget the real executor enforces. """ - if phase == "explore": - return int(min(CFG.horizon, CFG.max_num_steps_interaction_request)) + del phase return int(CFG.horizon) diff --git a/scripts/configs/empiric/common.yaml b/scripts/configs/empiric/common.yaml index ce700bf27c..9d8e2eef25 100644 --- a/scripts/configs/empiric/common.yaml +++ b/scripts/configs/empiric/common.yaml @@ -6,13 +6,9 @@ # Env FLAGS override approach FLAGS, which override these. ARGS: - debug -- make_failure_videos -- make_test_videos -- make_interaction_videos - auto_resume FLAGS: skill_phase_use_motion_planning: true - max_num_steps_interaction_request: 500 pretrained_model_service_provider: openrouter llm_model_name: google/gemini-2.5-pro llm_openai_max_response_tokens: 1e6 @@ -64,8 +60,3 @@ FLAGS: continual_belief_frame: true agent_model_repair: true continual_uncertainty_decisions: true - num_online_learning_cycles: 5 - online_learning_early_stopping: true - online_learning_early_stopping_require_all_attempts: true - online_learning_early_stopping_skip_redundant_test: true - online_nsrt_learning_requests_per_cycle: 2 diff --git a/scripts/configs/empiric/envs.yaml b/scripts/configs/empiric/envs.yaml index 7584e1bafd..0018744631 100644 --- a/scripts/configs/empiric/envs.yaml +++ b/scripts/configs/empiric/envs.yaml @@ -27,7 +27,6 @@ ENVS: FLAGS: max_initial_demos: 0 horizon: 3000 - max_num_steps_interaction_request: 2000 skill_place_settle_preload_force: 3.0 process_planning_heuristic_weight: 10.0 process_planning_max_execution_replans: 3 @@ -128,7 +127,6 @@ ENVS: domino_min_block_num_blues: 4 domino_block_cost: 0.1 domino_test_turn_ratio: 1.0 - online_learning_early_stopping_ignore_reward_bar: true continual_obs_noise_position: 0.01 continual_obs_noise_orientation: 0.04 continual_obs_noise_scalar: 0.0 diff --git a/tests/agent_sdk/prompt_goldens/learn_notes_message.md b/tests/agent_sdk/prompt_goldens/learn_notes_message.md deleted file mode 100644 index b93f415b47..0000000000 --- a/tests/agent_sdk/prompt_goldens/learn_notes_message.md +++ /dev/null @@ -1,35 +0,0 @@ -Write the world model document for this environment. There are 2 recorded trajectories (9 skill-level transitions) available: 0 oracle demonstration(s), which reached the goal by construction, and 2 interaction trajectory/ies collected during online learning, some of which may have failed to reach the goal. - - [0] interaction, task 0 - [1] interaction, task 0 - -Each trajectory carries a `train_task_idx`. `is_goal_state(state, task_idx)` (equivalently `train_tasks[task_idx].goal_holds(state)`) checks a single state for the goal atoms. Use it to confirm which trajectories reached the goal and to treat failed interaction trajectories as counterexamples: places where the environment disagreed with what a skill was expected to do. - -## Task goals (natural language) - -- Build the bridge. - -A `world_model.md` from an earlier cycle exists at `./world_model.md`. Read it first; this cycle's data may confirm, refine, or contradict what it says. Revise it in place. - -Data-structure source code is at: ./reference/structs.py - -## Available Predicates - -- Holding(robot:robot, block:block) - -## Object Types - -- robot: hand -- block: x, y, held - -## Options - -- Pick(robot:robot, block:block)[] - -## Available Tools - - - run_python - -## This session - -Read the data-structures file first, then explore the trajectory data with `run_python`. Write your world model to `./world_model.md` under the headings given in the system prompt, and finish with the deliverables listed there. diff --git a/tests/agent_sdk/prompt_goldens/learn_notes_system.md b/tests/agent_sdk/prompt_goldens/learn_notes_system.md deleted file mode 100644 index fe416ea3f5..0000000000 --- a/tests/agent_sdk/prompt_goldens/learn_notes_system.md +++ /dev/null @@ -1,31 +0,0 @@ -You are building a world model for a robotic manipulation environment as a natural-language document. No simulator will run from what you write: at planning time the same document is all the knowledge of the environment's dynamics the planner has, and it plans by reasoning over it, so what you write must let a careful reader predict what every skill does, when it works, and how the environment's own processes unfold over time. - -## What you produce - -One file, `world_model.md` (path given in the first message). Keep it organized under fixed headings so later cycles and the planner can find things: - -1. `# Mechanisms`: every process the environment runs on its own (delayed effects, gradual changes, propagation between objects, hidden state that changes what skills do), each with its trigger condition, its rate or duration in low-level steps, what it changes and by how much, and the evidence (trajectory and step) it comes from. -2. `# Skills`: for every skill, what it changes in the observed state when it succeeds (with the numbers: offsets, final poses, feature values as a function of the parameters), the conditions under which it fails and what the failure looks like, how many low-level steps it takes, and which of its continuous parameters matter and over what ranges. -3. `# Thresholds and geometry`: the quantitative gates the environment enforces (how close is close enough, which side of a fixture, what counts as supported), each bracketed by recorded attempts on both sides where the data allows. -4. `# Hidden state`: what the observation does not show, how it can be inferred from what it does show and from the history of skills executed, and how it evolves. -5. `# Recipes`: skill sequences, with parameter values, that the data shows reaching intermediate goals, and why they work. -6. `# Uncertainties and open questions`: what the data does not settle, phrased as the experiment that would settle it. - -Write for prediction, not description: a reader must be able to take a state and a skill call and write down the state after it. Prefer numbers over adjectives, and say where each number comes from. When you are unsure, say so and give the range. - -## Tools - -`run_python` is the one tool over the data: `trajectories` (`List[LowLevelTrajectory]`; each action's `get_option()` is the skill that produced it, so the skill-level transitions are the spans between skill changes), `describe_trajectory(i)`, `train_tasks`, `is_goal_state(state, task_idx)`, and `np`. Use `Read`, `Write` and `Edit` on `world_model.md`. - -## Deliverables of a learning session - -- The document, complete under the six headings above, with every mechanism the recorded episodes exercised reconciled against what you wrote before (earlier cycles' notes are yours to revise, not to append to). -- A decision record at the top: the key modeling commitments, the evidence behind each, and every hypothesis you kept without direct evidence, labelled as such. -- `./open_questions.md` with what the next exploration should collect first, and `./strategy.md` with how you would solve the train task given what you now know. - -## Workflow - -1. Explore the data with `run_python`: for each skill, which features change between its start and its end, under what conditions, and by how much; for each feature that changes while no skill touches it, what drives it. -2. `Write` or `Edit` `world_model.md`, one heading at a time, with the numbers and their evidence. -3. Check every claim against a transition it should predict: pick a recorded skill call, predict its outcome from your notes alone, compare. Fix the notes where the prediction is wrong. -4. Finish with the deliverables above. diff --git a/tests/agent_sdk/test_bilevel_sketch_regions.py b/tests/agent_sdk/test_bilevel_sketch_regions.py index 4c88070930..aed08f6b55 100644 --- a/tests/agent_sdk/test_bilevel_sketch_regions.py +++ b/tests/agent_sdk/test_bilevel_sketch_regions.py @@ -4,8 +4,7 @@ A region annotation gives a step's LLM-proposed params per-dimension half-widths: the exact center is tried once, then every later draw for the step is uniform inside ``clip([center - w, center + w], box)`` -instead of the full option box, taking precedence over any per-skill -sampler (region > sampler > uniform). +instead of the full option box. """ import asyncio @@ -334,19 +333,6 @@ def test_region_draws_confined_to_window(): assert 0.9 <= float(plan[0].params[0]) <= 0.95 -def test_region_takes_precedence_over_sampler(): - """A registered per-skill sampler is never consulted for a region step.""" - - def sampler(*_args): - raise AssertionError("sampler called despite region annotation") - - plan, success, _ = _refine(_region_step(0.5, 0.5), - max_samples_per_step=200, - parameterized_samplers={"Move": sampler}) - assert success - assert float(plan[0].params[0]) >= 0.9 - - def test_region_window_clipped_to_box(): """An oversized width clips to the option box (draws stay in-box).""" plan, success, _ = _refine(_region_step(0.95, 10.0), @@ -381,24 +367,6 @@ def test_region_applies_on_info_seeking_path(): assert float(plan[0].params[0]) >= 0.9 -def test_region_step_not_capped_by_deterministic_sampler(): - """A deterministic-flagged sampler must not collapse a region step to a - single attempt: the region bypasses the sampler entirely.""" - - def sampler(*_args): - return np.array([0.95], dtype=np.float32) - - sampler.deterministic = True - plan, success, total = _refine(_region_step(0.5, 0.5), - max_samples_per_step=200, - parameterized_samplers={"Move": sampler}) - assert success - # The failing center consumed the first attempt; regional draws (not a - # single deterministic try) then found a passing value. - assert total > 1 - assert float(plan[0].params[0]) >= 0.9 - - def _named_step(fn, name="named"): return SketchStep(option=_Move, objects=[_block], diff --git a/tests/agent_sdk/test_bilevel_sketch_samplers.py b/tests/agent_sdk/test_bilevel_sketch_samplers.py index 93c58a8596..ca350e2bf5 100644 --- a/tests/agent_sdk/test_bilevel_sketch_samplers.py +++ b/tests/agent_sdk/test_bilevel_sketch_samplers.py @@ -1,12 +1,8 @@ -"""Tests for per-skill synthesized samplers in ``sketch_refinement``. - -Verifies that a sampler registered under an option name in -``parameterized_samplers`` is consulted (with the step's subgoal + -objects + the option's params box) to draw that option's continuous -params during refinement — on both the plain and info-seeking paths — -and that a missing / misbehaving sampler falls back to uniform sampling -so refinement is byte-for-byte unchanged when no usable sampler is -supplied. +"""Tests for how ``sketch_refinement`` draws a step's continuous params: + +uniformly by default, the agent's proposed params first, and on the +info-seeking path a pool of the proposal and sampled draws; plus the +sketch parsing and forward-execution helpers. """ # pylint: disable=unused-import @@ -23,7 +19,7 @@ parse_sketch_from_text, strip_subgoal_annotations from predicators.agent_sdk.sketch_refinement import \ refine_and_validate_report, refine_sketch, sample_params -from predicators.agent_sdk.sketch_types import SketchStep +from predicators.agent_sdk.sketch_types import GroundSampler, SketchStep from predicators.structs import Action, GroundAtom, Object, \ ParameterizedOption, Predicate, State, Task, Type @@ -126,107 +122,8 @@ def _easy_task_and_sketch(): return task, sketch -def test_registered_sampler_is_used(): - """A targeted sampler lands the hard subgoal on the first sample.""" - calls = [] - - def sampler(state, subgoal_atoms, rng, objects): - del state, rng - calls.append((objects, subgoal_atoms)) - return np.array([0.95], dtype=np.float32) - - model = _FakeOptionModel() - plan, success, total = refine_sketch( - _task_hi(), - _sketch_hi(), - model, - predicates={_ReachedHi}, - timeout=10.0, - rng=np.random.default_rng(0), - max_samples_per_step=50, - check_subgoals=True, - check_final_goal=False, - parameterized_samplers={"Move": sampler}) - assert success - assert np.isclose(float(plan[0].params[0]), 0.95) - # Feasible on the very first attempt — none of the uniform churn. - assert total == 1 - assert model.num_calls == 1 - # The sampler saw the right subgoal and objects. - objs, subgoal = calls[0] - assert [o.name for o in objs] == ["block0"] - assert GroundAtom(_ReachedHi, [_block]) in subgoal - - -def test_missing_entry_falls_back_to_uniform(): - """A sampler keyed by another option leaves Move on the uniform path.""" - seed = 7 - first = float(sample_params(_Move, np.random.default_rng(seed))[0]) - task, sketch = _easy_task_and_sketch() - - def other(*_args): - raise AssertionError("sampler for a different option was called") - - plan, success, _ = refine_sketch( - task, - sketch, - _FakeOptionModel(), - predicates={_Reached}, - timeout=10.0, - rng=np.random.default_rng(seed), - max_samples_per_step=50, - check_subgoals=True, - check_final_goal=False, - parameterized_samplers={"OtherOption": other}) - assert success - # Identical to the no-sampler uniform draw. - assert float(plan[0].params[0]) == first - - -def test_bad_shape_falls_back_to_uniform(): - """A wrong-shaped return is rejected; uniform sampling still succeeds.""" - task, sketch = _easy_task_and_sketch() - - def bad(*_args): - return np.array([0.5, 0.5], dtype=np.float32) # shape (2,) != (1,) - - plan, success, _ = refine_sketch(task, - sketch, - _FakeOptionModel(), - predicates={_Reached}, - timeout=10.0, - rng=np.random.default_rng(0), - max_samples_per_step=50, - check_subgoals=True, - check_final_goal=False, - parameterized_samplers={"Move": bad}) - assert success - assert 0.0 <= float(plan[0].params[0]) <= 1.0 - - -def test_raising_sampler_falls_back_to_uniform(): - """A sampler that raises is caught and uniform sampling proceeds.""" - task, sketch = _easy_task_and_sketch() - - def boom(*_args): - raise ValueError("nope") - - _, success, _ = refine_sketch(task, - sketch, - _FakeOptionModel(), - predicates={_Reached}, - timeout=10.0, - rng=np.random.default_rng(0), - max_samples_per_step=50, - check_subgoals=True, - check_final_goal=False, - parameterized_samplers={"Move": boom}) - assert success - - -def test_none_samplers_unchanged(): - """parameterized_samplers=None reproduces the plain first-uniform-draw - param.""" +def test_plain_step_draws_uniformly(): + """A step with no ground sampler takes the plain first uniform draw.""" seed = 7 first = float(sample_params(_Move, np.random.default_rng(seed))[0]) task, sketch = _easy_task_and_sketch() @@ -238,38 +135,11 @@ def test_none_samplers_unchanged(): rng=np.random.default_rng(seed), max_samples_per_step=50, check_subgoals=True, - check_final_goal=False, - parameterized_samplers=None) + check_final_goal=False) assert success assert float(plan[0].params[0]) == first -def test_sampler_used_on_info_seeking_path(): - """The info-seeking draw loop also routes through the sampler.""" - - def sampler(_s, _a, rng, _o): - # Jitter so candidates differ but all clear the x>=0.9 subgoal. - return np.array([0.9 + 0.05 * rng.random()], dtype=np.float32) - - model = _FakeOptionModel() - plan, success, _ = refine_sketch( - _task_hi(), - _sketch_hi(), - model, - predicates={_ReachedHi}, - timeout=10.0, - rng=np.random.default_rng(0), - max_samples_per_step=50, - check_subgoals=True, - check_final_goal=False, - info_scorer=lambda s, _a: s.get(_block, "x"), - info_n_feasible_target=4, - parameterized_samplers={"Move": sampler}) - assert success - # Every pooled candidate came from the sampler => satisfies x >= 0.9. - assert float(plan[0].params[0]) >= 0.9 - - # --------------------------------------------------------------------------- # # LLM-proposed initial_params (tried first, with sampling fallback). # --------------------------------------------------------------------------- # @@ -289,8 +159,7 @@ def test_initial_params_tried_first_without_sampler(): rng=np.random.default_rng(0), max_samples_per_step=50, check_subgoals=True, - check_final_goal=False, - parameterized_samplers=None) + check_final_goal=False) assert success # The proposal satisfied the hard subgoal on the very first attempt. assert np.isclose(float(plan[0].params[0]), 0.95) @@ -312,8 +181,7 @@ def test_initial_params_fall_back_to_uniform_on_failure(): rng=np.random.default_rng(0), max_samples_per_step=200, check_subgoals=True, - check_final_goal=False, - parameterized_samplers=None) + check_final_goal=False) assert success # The failed proposal was the first sample; uniform then found x >= 0.9. assert total > 1 @@ -333,14 +201,18 @@ def test_initial_params_clipped_to_box(): rng=np.random.default_rng(0), max_samples_per_step=50, check_subgoals=True, - check_final_goal=False, - parameterized_samplers=None) + check_final_goal=False) assert success # 5.0 clipped to the option's high bound (1.0), which clears x >= 0.9. assert np.isclose(float(plan[0].params[0]), 1.0) assert total == 1 +def _near_hi(_s, _a, rng, _o): + """Sampled candidates that clear x >= 0.9 but stay below 1.0.""" + return np.array([0.9 + 0.05 * rng.random()], dtype=np.float32) + + def test_initial_params_seeded_and_win_on_disagreement(): """LLM params are pooled with sampled draws; the argmax (most informative) is chosen. @@ -350,13 +222,11 @@ def test_initial_params_seeded_and_win_on_disagreement(): step = SketchStep(option=_Move, objects=[_block], subgoal_atoms={GroundAtom(_ReachedHi, [_block])}, - initial_params=np.array([1.0], dtype=np.float32)) + initial_params=np.array([1.0], dtype=np.float32), + ground_sampler=GroundSampler(fn=_near_hi, + name="near_hi")) model = _FakeOptionModel() - # Sampled candidates clear x >= 0.9 but stay below the guess's x = 1.0. - def sampler(_s, _a, rng, _o): - return np.array([0.9 + 0.05 * rng.random()], dtype=np.float32) - plan, success, _ = refine_sketch( _task_hi(), [step], model, @@ -367,8 +237,7 @@ def sampler(_s, _a, rng, _o): check_subgoals=True, check_final_goal=False, info_scorer=lambda s, _a: float(s.get(_block, "x")), - info_n_feasible_target=4, - parameterized_samplers={"Move": sampler}) + info_n_feasible_target=4) assert success # The guess had the highest disagreement (x = 1.0) => argmax picked it. assert np.isclose(float(plan[0].params[0]), 1.0) @@ -383,7 +252,10 @@ def test_initial_params_lose_to_more_informative_draw(): step = SketchStep(option=_Move, objects=[_block], subgoal_atoms={GroundAtom(_ReachedHi, [_block])}, - initial_params=np.array([0.9], dtype=np.float32)) + initial_params=np.array([0.9], dtype=np.float32), + ground_sampler=GroundSampler( + fn=lambda *_a: np.array([0.99], dtype=np.float32), + name="at_0_99")) plan, success, _ = refine_sketch( _task_hi(), [step], _FakeOptionModel(), @@ -394,10 +266,7 @@ def test_initial_params_lose_to_more_informative_draw(): check_subgoals=True, check_final_goal=False, info_scorer=lambda s, _a: float(s.get(_block, "x")), - info_n_feasible_target=4, - parameterized_samplers={ - "Move": lambda *_a: np.array([0.99], dtype=np.float32) - }) + info_n_feasible_target=4) assert success # The seeded guess (x = 0.9) was beaten by the more informative draw # (0.99) — proving it is pooled, not accepted just for being first. @@ -409,7 +278,9 @@ def test_initial_params_infeasible_seed_info_seeking_recovers(): step = SketchStep(option=_Move, objects=[_block], subgoal_atoms={GroundAtom(_ReachedHi, [_block])}, - initial_params=np.array([0.0], dtype=np.float32)) + initial_params=np.array([0.0], dtype=np.float32), + ground_sampler=GroundSampler(fn=_near_hi, + name="near_hi")) plan, success, _ = refine_sketch( _task_hi(), [step], _FakeOptionModel(), @@ -420,12 +291,7 @@ def test_initial_params_infeasible_seed_info_seeking_recovers(): check_subgoals=True, check_final_goal=False, info_scorer=lambda s, _a: float(s.get(_block, "x")), - info_n_feasible_target=4, - parameterized_samplers={ - "Move": - lambda _s, _a, rng, _o: np.array([0.9 + 0.05 * rng.random()], - dtype=np.float32) - }) + info_n_feasible_target=4) assert success # The infeasible guess (x = 0) wasn't pooled; sampled candidates won. assert float(plan[0].params[0]) >= 0.9 diff --git a/tests/agent_sdk/test_docker_agent_runner.py b/tests/agent_sdk/test_docker_agent_runner.py deleted file mode 100644 index 0cdac4e306..0000000000 --- a/tests/agent_sdk/test_docker_agent_runner.py +++ /dev/null @@ -1,165 +0,0 @@ -"""Tests for stale Object hash repair after cross-process unpickling. - -Characterizes ``_rehash_objects_after_unpickle``: Object caches -``_hash = hash(str(self))`` in a cached_property, and PYTHONHASHSEED -randomization makes those cached values stale in a new process. The -tests simulate that by planting a bogus ``_hash`` before building the -``State.data`` dict, so the dict's stored entry hashes are stale, which -is exactly what a cross-seed dill roundtrip produces (verified E2E). - -Characterized behavior (current, not aspirational): - -- The rehash clears every reachable Object's cached ``_hash``/``_str``, - so stale and freshly created equal Objects hash identically again. -- The ``state.data = dict(state.data)`` rebuild does NOT re-key stale - entries: CPython's ``dict(d)`` copies each entry's stored hash without - recomputing it, so lookups into pre-existing dicts remain broken (a - dict comprehension would repair them). Tests below pin this down. - -Only ``_rehash_objects_after_unpickle`` is under test; ``main()`` and -``_run_query`` need Docker and the SDK. -""" -# pylint: disable=protected-access -from types import SimpleNamespace - -import numpy as np -import pytest - -from predicators.agent_sdk.docker_agent_runner import \ - _rehash_objects_after_unpickle -from predicators.structs import Action, GroundAtom, LowLevelTrajectory, \ - Object, Predicate, State, Task, Type - -_block_type = Type("block", ["x"]) -_OnTable = Predicate("OnTable", [_block_type], lambda s, o: True) - - -def _make_state(obj, value=0.0): - return State({obj: np.array([value], dtype=np.float32)}) - - -def _corrupt_hash(obj): - """Plant a stale cached hash, as if pickled under another hash seed.""" - true_hash = hash(obj) # Populate the cached_property. - # +1 guarantees a different bucket index modulo any power-of-two - # table size, so corrupted lookups miss deterministically. - obj.__dict__["_hash"] = true_hash + 1 - return true_hash - - -def _stale_state(name="block0"): - """Build a state whose dict entries are stored under a stale hash. - - Corrupting before insertion mirrors unpickling: dill inserts keys - while their cached (old-process) ``_hash`` is in effect. - """ - obj = Object(name, _block_type) - _corrupt_hash(obj) - state = _make_state(obj) - return obj, state - - -def test_stale_hash_breaks_fresh_object_lookup(): - """Precondition: a fresh equal Object cannot find the stale key.""" - stale_obj, state = _stale_state() - fresh_obj = Object("block0", _block_type) - assert fresh_obj == stale_obj - assert hash(fresh_obj) != hash(stale_obj) - with pytest.raises(KeyError): - _ = state.data[fresh_obj] - # The stale instance itself still works: its cached hash matches - # the hash stored at insertion. - assert state.data[stale_obj] is not None - - -def test_rehash_clears_hash_caches_on_all_ctx_surfaces(): - """All reachable Objects re-hash equal to fresh ones after rehash.""" - train_obj, train_state = _stale_state("train_block") - train_task = Task(train_state, {GroundAtom(_OnTable, [train_obj])}) - - cur_obj, cur_state = _stale_state("cur_block") - cur_task = Task(cur_state, {GroundAtom(_OnTable, [cur_obj])}) - - example_obj, example_state = _stale_state("example_block") - - traj_obj, traj_state0 = _stale_state("traj_block") - # traj_obj still carries the stale cache, so this second dict is - # keyed under it too, like a second unpickled state. - traj_state1 = _make_state(traj_obj, value=1.0) - traj = LowLevelTrajectory([traj_state0, traj_state1], - [Action(np.zeros(1, dtype=np.float32))]) - - ctx = SimpleNamespace(train_tasks=[train_task], - current_task=cur_task, - example_state=example_state, - offline_trajectories=[traj], - online_trajectories=[]) - _rehash_objects_after_unpickle(ctx) - - for name, obj in (("train_block", train_obj), ("cur_block", cur_obj), - ("example_block", example_obj), ("traj_block", - traj_obj)): - fresh = Object(name, _block_type) - assert hash(obj) == hash(fresh), f"cache not cleared for {name}" - # Goal-atom objects were cleared too (same instances here, but the - # atom path is walked independently of the state path). - goal_obj = next(iter(cur_task.goal)).objects[0] - assert hash(goal_obj) == hash(Object("cur_block", _block_type)) - - -def test_rehash_rekeys_stale_entries(): - """The rebuild re-keys entries under the repaired hashes. - - The comprehension in ``_process_state`` is load-bearing: ``dict(d)`` - would copy each entry's stored hash without calling ``__hash__``, - leaving lookups broken even after the caches are cleared. - """ - stale_obj, state = _stale_state() - ctx = SimpleNamespace(train_tasks=[Task(state, set())], current_task=None) - _rehash_objects_after_unpickle(ctx) - - fresh_obj = Object("block0", _block_type) - assert hash(stale_obj) == hash(fresh_obj) # Caches repaired... - assert fresh_obj in state.data # ...and the table re-keyed. - assert state.data[fresh_obj][0] == pytest.approx(0.0) - assert stale_obj in state.data - - -def test_rehash_makes_newly_built_dicts_consistent(): - """Dicts re-keyed by hand after the rehash serve fresh lookups.""" - stale_obj, state = _stale_state() - ctx = SimpleNamespace(train_tasks=[Task(state, set())], current_task=None) - _rehash_objects_after_unpickle(ctx) - # Re-inserting under the repaired hashes (what a comprehension or - # any post-rehash construction does) restores fresh-object lookups. - # pylint: disable-next=unnecessary-comprehension - rekeyed = {obj: val for obj, val in state.data.items()} - fresh_obj = Object("block0", _block_type) - assert fresh_obj in rekeyed - assert rekeyed[fresh_obj][0] == pytest.approx(0.0) - assert stale_obj in rekeyed - - -def test_rehash_processes_environment_task_init_obs(): - """Task-likes exposing init_obs (EnvironmentTask shape) are walked.""" - stale_obj, state = _stale_state("obs_block") - env_task = SimpleNamespace(init_obs=state, goal_description=None) - ctx = SimpleNamespace(train_tasks=[env_task], current_task=None) - _rehash_objects_after_unpickle(ctx) - assert hash(stale_obj) == hash(Object("obs_block", _block_type)) - - -def test_rehash_handles_minimal_ctx_and_none_current_task(): - """A ctx with no tasks, states, or trajectories is a no-op.""" - ctx = SimpleNamespace(current_task=None) - _rehash_objects_after_unpickle(ctx) # Should not raise. - - -def test_rehash_is_idempotent_on_hash_caches(): - """Running the rehash twice keeps hashes consistent.""" - stale_obj, state = _stale_state() - task = Task(state, set()) - ctx = SimpleNamespace(train_tasks=[task], current_task=task) - _rehash_objects_after_unpickle(ctx) - _rehash_objects_after_unpickle(ctx) - assert hash(stale_obj) == hash(Object("block0", _block_type)) diff --git a/tests/agent_sdk/test_ground_sampler_loader.py b/tests/agent_sdk/test_ground_sampler_loader.py new file mode 100644 index 0000000000..7031fba846 --- /dev/null +++ b/tests/agent_sdk/test_ground_sampler_loader.py @@ -0,0 +1,61 @@ +"""Tests for loading the agent's ``GROUND_SAMPLERS``, the named samplers a +sketch step references with ``~ name``.""" + +import numpy as np +from gym.spaces import Box + +from predicators.agent_sdk.proposal_exec import build_exec_context, \ + load_ground_samplers +from predicators.structs import Action, Object, ParameterizedOption, \ + Predicate, Type + +_block_type = Type("block", ["x"]) +_block = Object("block0", _block_type) + +_Reached = Predicate("Reached", [_block_type], lambda s, o: True) + +_Move = ParameterizedOption( + "Move", + types=[_block_type], + params_space=Box(low=np.array([0.0], dtype=np.float32), + high=np.array([1.0], dtype=np.float32)), + policy=lambda _s, _m, _o, _p: Action(np.zeros(1, dtype=np.float32)), + initiable=lambda _s, _m, _o, _p: True, + terminal=lambda _s, _m, _o, _p: False, +) + + +def test_load_ground_samplers_happy_and_bad_entries(): + """GROUND_SAMPLERS loads callables; bad keys/values warn and drop.""" + ctx = build_exec_context(types={_block_type}, + predicates={_Reached}, + options={_Move}) + code = """\ +def _fn(state, subgoal_atoms, rng, objects): + del state, subgoal_atoms, objects + return np.array([0.5], dtype=np.float32) + +GROUND_SAMPLERS = {"hi_band": _fn, "not-an-identifier": _fn, "seven": 7} +""" + fns, warnings, err = load_ground_samplers(code, ctx) + assert err is None + assert set(fns) == {"hi_band"} + assert len(warnings) == 2 + assert any("identifiers" in w for w in warnings) + assert any("not callable" in w for w in warnings) + + +def test_load_ground_samplers_errors(): + """Exec failures and non-dict bindings load nothing, with an error.""" + ctx = build_exec_context(types={_block_type}, + predicates={_Reached}, + options={_Move}) + fns, _, err = load_ground_samplers("raise RuntimeError('boom')", ctx) + assert not fns + assert err is not None and "boom" in err + ctx = build_exec_context(types={_block_type}, + predicates={_Reached}, + options={_Move}) + fns, _, err = load_ground_samplers("GROUND_SAMPLERS = [1]", ctx) + assert not fns + assert err is not None and "must be a dict" in err diff --git a/tests/agent_sdk/test_probe_synthesis.py b/tests/agent_sdk/test_probe_synthesis.py index b30af7d3e8..309e75c62e 100644 --- a/tests/agent_sdk/test_probe_synthesis.py +++ b/tests/agent_sdk/test_probe_synthesis.py @@ -20,7 +20,6 @@ create_mcp_tools from predicators.approaches.agent_sim_learning_approach import \ AgentSimLearningApproach -from predicators.code_sim_learning.fit_space import FitResult from predicators.option_model import _OptionModelBase from predicators.structs import Object, State, Task, Type @@ -83,7 +82,7 @@ def test_probe_reset_requires_task_idx_during_synthesis() -> None: assert sim._state is not None -def test_candidate_probe_model_provider_glue(tmp_path, monkeypatch) -> None: +def test_candidate_probe_model_provider_glue(tmp_path) -> None: """The provider gates on a loadable simulator.py, caches by content hash, rebuilds on change, and NEVER fits: the candidate runs at carried-over. @@ -92,36 +91,19 @@ def test_candidate_probe_model_provider_glue(tmp_path, monkeypatch) -> None: runs at the fitted values (status fitted). Exercises the real ``_make_candidate_probe_model_provider`` and the - real file loader; only the fit/build layer below - ``build_candidate_option_model`` is stubbed (its body is the shared - ``evaluate_plan_refinement`` path). + real file loader; only the option-model build layer below + ``build_candidate_option_model`` is stubbed. """ approach = object.__new__(AgentSimLearningApproach) approach._fitted_params = {} approach._latent_init = None approach._tool_context = ToolContext() - fit_calls = {"n": 0} - - def _fake_fit(rules, specs, triples, features): - del rules, triples, features - fit_calls["n"] += 1 - names = [s.name for s in specs] - return FitResult(names=names, - samples=np.array([[s.init_value for s in specs]]), - log_probs=np.array([0.0])), 0.0 - - monkeypatch.setattr( - "predicators.approaches.synthesis_validation.fit_rule_parameters", - _fake_fit) setattr(approach, "_build_combined_simulator", lambda learned: learned) setattr(approach, "_build_option_model", lambda sim: ("model", sim)) simulator_file = str(tmp_path / "simulator.py") - provider = approach._make_candidate_probe_model_provider( - simulator_file, - trajectories=[], - base_pred_triples=[], - inferred_hint={"thing": ["x"]}) + provider = approach._make_candidate_probe_model_provider(simulator_file, + trajectories=[]) # No file yet: hard error, never a fallback model. with pytest.raises(RuntimeError, match="no candidate simulator yet"): @@ -142,7 +124,6 @@ def _fake_fit(rules, specs, triples, features): f.write(valid) model = provider() # No implicit fit: declared init value, and the status says so. - assert fit_calls["n"] == 0 assert approach._fitted_params == {"k": 1.0} status = approach._tool_context.probe_param_status assert status is not None and status.startswith("UNFITTED") @@ -159,7 +140,6 @@ def _fake_fit(rules, specs, triples, features): assert approach._fitted_params == {"k": 1.7} assert approach._tool_context.probe_param_status == \ "fitted (cycle_000_vers_002)" - assert fit_calls["n"] == 0 # A rejected fit remains unvalidated through rebuild and cache reuse. approach._publish_probe_fit({"k": 1.0}, @@ -179,7 +159,6 @@ def _fake_fit(rules, specs, triples, features): provider() assert approach._tool_context.probe_param_status.startswith("PARTIAL FIT") assert "2/3" in approach._tool_context.probe_param_status - assert fit_calls["n"] == 0 # Changed content: rebuilt UNFITTED, carrying the last fit's value # for a param that still exists inside its box. diff --git a/tests/agent_sdk/test_prompt_goldens.py b/tests/agent_sdk/test_prompt_goldens.py index a04b765e4a..774a55c8dd 100644 --- a/tests/agent_sdk/test_prompt_goldens.py +++ b/tests/agent_sdk/test_prompt_goldens.py @@ -373,34 +373,6 @@ def test_golden_learn_program_message() -> None: ])) -def test_golden_learn_notes_system() -> None: - """The natural-language world-model learn system prompt (paper arm C3).""" - _check_golden("learn_notes_system", - learn_prompts.build_notes_learn_system_prompt()) - - -def test_golden_learn_notes_message() -> None: - """The natural-language world-model learn first message with prior notes - and a goal.""" - _check_golden( - "learn_notes_message", - learn_prompts.build_notes_learn_message( - n_trajs=2, - n_transitions=9, - n_demos=0, - n_interaction=2, - trajectory_listing=" [0] interaction, task 0\n" - " [1] interaction, task 0", - structs_ref="./reference/structs.py", - predicate_listing="- Holding(robot:robot, block:block)", - types_digest="- robot: hand\n- block: x, y, held", - options_digest="- Pick(robot:robot, block:block)[]", - notes_file="./world_model.md", - goal_nls=["Build the bridge.", "Build the bridge."], - has_prior_notes=True, - tools_block=learn_prompts.render_tools_block(["run_python"]))) - - @pytest.mark.parametrize("model_based,noise,repair", [ (True, False, False), (True, True, True), diff --git a/tests/agent_sdk/test_sampler_synthesis_tools.py b/tests/agent_sdk/test_sampler_synthesis_tools.py deleted file mode 100644 index 9ab8f50ded..0000000000 --- a/tests/agent_sdk/test_sampler_synthesis_tools.py +++ /dev/null @@ -1,241 +0,0 @@ -"""Tests for the ``sim.samplers()`` loader (make_sampler_loader). - -Drives the real tool handler against a stub approach: loading -``LEARNED_SAMPLERS`` from ``samplers.py``, installing the validated dict -onto the approach, skip warnings for bad entries, error paths, snapshot -versioning, and the sanity check's empty-subgoal-set contract. -""" - -from typing import Any, Dict, Set - -import numpy as np -from gym.spaces import Box - -from predicators import utils -from predicators.agent_sdk.belief_probe import BeliefProbe -from predicators.agent_sdk.proposal_exec import build_exec_context, \ - load_ground_samplers -from predicators.agent_sdk.tools import ToolContext, make_sampler_loader -from predicators.structs import Action, GroundAtom, Object, \ - ParameterizedOption, Predicate, State, Task, Type - -_block_type = Type("block", ["x"]) -_block = Object("block0", _block_type) - -_Reached = Predicate("Reached", [_block_type], lambda s, o: True) - -_Move = ParameterizedOption( - "Move", - types=[_block_type], - params_space=Box(low=np.array([0.0], dtype=np.float32), - high=np.array([1.0], dtype=np.float32)), - policy=lambda _s, _m, _o, _p: Action(np.zeros(1, dtype=np.float32)), - initiable=lambda _s, _m, _o, _p: True, - terminal=lambda _s, _m, _o, _p: False, -) - - -class _StubApproach: - """The minimal approach surface make_sampler_loader uses.""" - - def __init__(self): - init = State({_block: np.array([0.0], dtype=np.float32)}) - self._types = {_block_type} - self._train_tasks = [Task(init, {GroundAtom(_Reached, [_block])})] - self._fitted_params: Dict[str, float] = {} - self._synthesized_samplers: Dict[str, Any] = {} - - def _get_all_predicates(self) -> Set[Predicate]: - return {_Reached} - - def _get_all_options(self) -> Set[ParameterizedOption]: - return {_Move} - - -def _run_sampler_loader(tmp_path, code=None): - utils.reset_config({"seed": 0}) - samplers_file = str(tmp_path / "samplers.py") - if code is not None: - with open(samplers_file, "w", encoding="utf-8") as f: - f.write(code) - approach = _StubApproach() - loader = make_sampler_loader( - samplers_file=samplers_file, - samplers_versions_dir=str(tmp_path / "samplers_versions"), - approach=approach, - cycle_index_provider=lambda: 1, - ) - # Through the probe, exactly as the agent reaches it. - ctx = ToolContext() - ctx.probe_artifact_loaders["samplers"] = loader - return BeliefProbe(ctx).samplers(), approach - - -_GOOD = """\ -def _move_sampler(state, subgoal_atoms, rng, objects): - del state, subgoal_atoms, objects - return np.array([0.25 + 0.01 * rng.random()], dtype=np.float32) - -LEARNED_SAMPLERS = {"Move": _move_sampler} -""" - - -def test_evaluate_sampler_installs_valid_samplers(tmp_path): - """A valid samplers.py is installed onto the approach and passes the sanity - check.""" - text, approach = _run_sampler_loader(tmp_path, _GOOD) - assert "1 per-skill sampler(s) installed" in text - assert "Move: OK" in text - assert "3/3 within the params box" in text - assert set(approach._synthesized_samplers) == {"Move"} # pylint: disable=protected-access - assert callable(approach._synthesized_samplers["Move"]) # pylint: disable=protected-access - - -def test_evaluate_sampler_warns_unknown_option(tmp_path): - """An entry keyed by a non-option name is skipped with a warning.""" - code = _GOOD + "\nLEARNED_SAMPLERS['Teleport'] = _move_sampler\n" - text, approach = _run_sampler_loader(tmp_path, code) - assert "Skipped 'Teleport' (not a known option name" in text - assert set(approach._synthesized_samplers) == {"Move"} # pylint: disable=protected-access - - -def test_evaluate_sampler_warns_non_callable(tmp_path): - """A non-callable value is skipped with a warning.""" - code = _GOOD + "\nLEARNED_SAMPLERS['Move'] = 3.0\n" - text, approach = _run_sampler_loader(tmp_path, code) - assert "Skipped 'Move' (value is not callable" in text - assert not approach._synthesized_samplers # pylint: disable=protected-access - - -def test_evaluate_sampler_reports_exec_error(tmp_path): - """A samplers.py that raises at import time reports the traceback.""" - text, approach = _run_sampler_loader(tmp_path, - "raise RuntimeError('boom')") - assert "Error executing" in text - assert "boom" in text - assert not approach._synthesized_samplers # pylint: disable=protected-access - - -def test_evaluate_sampler_reports_missing_symbol(tmp_path): - """A file without LEARNED_SAMPLERS names the missing symbol.""" - text, _ = _run_sampler_loader(tmp_path, "x = 1\n") - assert "LEARNED_SAMPLERS" in text - - -def test_evaluate_sampler_missing_file_hint(tmp_path): - """A missing samplers.py returns the Write hint, not a crash.""" - text, _ = _run_sampler_loader(tmp_path, code=None) - assert "Use Write to create it" in text - - -def test_evaluate_sampler_empty_dict_message(tmp_path): - """An empty LEARNED_SAMPLERS asks for entries instead of sanity lines.""" - text, _ = _run_sampler_loader(tmp_path, "LEARNED_SAMPLERS = {}\n") - assert "LEARNED_SAMPLERS is empty" in text - assert "Sanity check" not in text - - -def test_evaluate_sampler_version_tag_bumps_on_edit(tmp_path): - """Within one tool instance (one synthesis session), an edited samplers.py - gets a fresh version tag; an identical reload keeps it.""" - utils.reset_config({"seed": 0}) - samplers_file = tmp_path / "samplers.py" - samplers_file.write_text(_GOOD, encoding="utf-8") - loader = make_sampler_loader( - samplers_file=str(samplers_file), - samplers_versions_dir=str(tmp_path / "samplers_versions"), - approach=_StubApproach(), - cycle_index_provider=lambda: 1, - ) - - def _tag(): - return loader().split("]")[0].lstrip("[") - - tag1 = _tag() - tag2 = _tag() # unchanged file: same tag (snapshot deduped) - samplers_file.write_text(_GOOD.replace("0.25", "0.75"), encoding="utf-8") - tag3 = _tag() - assert tag1 == "cycle_001_vers_001" - assert tag2 == tag1 - assert tag3 == "cycle_001_vers_002" - - -def test_sanity_check_raising_sampler_mentions_empty_subgoal_contract( - tmp_path): - """A sampler that assumes a non-empty subgoal set gets the contract spelled - out: refinement calls samplers with subgoal_atoms=set() at steps with no - annotation, and the sanity check does the same.""" - code = """\ -def _needs_subgoal(state, subgoal_atoms, rng, objects): - atom = next(iter(subgoal_atoms)) - del state, rng, objects, atom - return np.array([0.5], dtype=np.float32) - -LEARNED_SAMPLERS = {"Move": _needs_subgoal} -""" - text, _ = _run_sampler_loader(tmp_path, code) - assert "ERROR" in text - assert "subgoal_atoms=set()" in text - assert "must not crash on an empty set" in text - - -def test_load_ground_samplers_happy_and_bad_entries(): - """GROUND_SAMPLERS loads callables; bad keys/values warn and drop.""" - ctx = build_exec_context(types={_block_type}, - predicates={_Reached}, - options={_Move}) - code = """\ -def _fn(state, subgoal_atoms, rng, objects): - del state, subgoal_atoms, objects - return np.array([0.5], dtype=np.float32) - -GROUND_SAMPLERS = {"hi_band": _fn, "not-an-identifier": _fn, "seven": 7} -""" - fns, warnings, err = load_ground_samplers(code, ctx) - assert err is None - assert set(fns) == {"hi_band"} - assert len(warnings) == 2 - assert any("identifiers" in w for w in warnings) - assert any("not callable" in w for w in warnings) - - -def test_load_ground_samplers_errors(): - """Exec failures and non-dict bindings load nothing, with an error.""" - ctx = build_exec_context(types={_block_type}, - predicates={_Reached}, - options={_Move}) - fns, _, err = load_ground_samplers("raise RuntimeError('boom')", ctx) - assert not fns - assert err is not None and "boom" in err - ctx = build_exec_context(types={_block_type}, - predicates={_Reached}, - options={_Move}) - fns, _, err = load_ground_samplers("GROUND_SAMPLERS = [1]", ctx) - assert not fns - assert err is not None and "must be a dict" in err - - -def test_sanity_check_wrong_shape_reports_error(tmp_path): - """A wrong-shaped return is reported with got/expected shapes.""" - code = """\ -def _bad_shape(state, subgoal_atoms, rng, objects): - del state, subgoal_atoms, rng, objects - return np.array([0.5, 0.5], dtype=np.float32) - -LEARNED_SAMPLERS = {"Move": _bad_shape} -""" - text, _ = _run_sampler_loader(tmp_path, code) - assert "ERROR" in text - assert "returned shape (2,), expected (1,)" in text - - -def test_probe_samplers_unavailable_without_a_loader(): - """Outside a sampler-synthesis session the probe has no samplers.py surface - and says so instead of silently doing nothing.""" - probe = BeliefProbe(ToolContext()) - try: - probe.samplers() - except RuntimeError as e: - assert "sim.samplers is unavailable" in str(e) - else: - raise AssertionError("sim.samplers() must raise without a loader") diff --git a/tests/agent_sdk/test_submit_plan_capture.py b/tests/agent_sdk/test_submit_plan_capture.py index 12fa29558e..a4b0babcce 100644 --- a/tests/agent_sdk/test_submit_plan_capture.py +++ b/tests/agent_sdk/test_submit_plan_capture.py @@ -1105,30 +1105,6 @@ def test_latent_only_annotation_excluded_from_monitoring(): {"ReachedHi(block0:block)"} -def test_raising_negative_annotation_excluded_from_monitoring(): - """A negative annotation whose classifier RAISES without a latent is. - - dropped too - the monitor could not evaluate it on a real state. - """ - # Empty bond set: RaisingBond is False with the latent, so the NOT - # annotation holds in the belief rollout and survives to the probe. - model = _LatentModel(set()) - utils.reset_config({"agent_plan_validation_rollouts": 3}) - ctx = _make_ctx(model) - ctx.predicates.add(_RaisingBond) - # Belief rollouts attach an initial latent to the task init (the - # production path does this via _attach_initial_latent); without it - # the classifier would raise inside the rollout itself. - ctx.current_task.init.latent = {"_bonds": set()} - text = _call_tool(ctx, _NEG_RAISING_PLAN) - assert "Captured as the current answer" in text - assert "cannot be verified from a real observation" in text - assert "RaisingBond" in text - sketch = ctx.solved_sketch - assert sketch is not None - assert not sketch[0].subgoal_neg_atoms - - def test_observation_backed_annotation_survives_probe(): """An annotation that holds from observable features alone is kept: diff --git a/tests/agent_sdk/test_tool_registry.py b/tests/agent_sdk/test_tool_registry.py index 55c6963402..9d67423531 100644 --- a/tests/agent_sdk/test_tool_registry.py +++ b/tests/agent_sdk/test_tool_registry.py @@ -312,9 +312,8 @@ def test_agent_render_resolution() -> None: def test_synthesis_tool_names_run_python() -> None: - """Every session offers one ``run_python``: the synthesis roster carries - its own instance (fit data + the candidate-simulator probe in one - namespace), the solve roster the probe over the deployed belief model. + """The synthesis roster carries one ``run_python``, its own instance (fit + data + the candidate-simulator probe in one namespace). Fitting, residual reports, plan validation, and scene work are probe methods (``sim.fit`` / ``sim.residuals`` / ``sim.refine`` / @@ -328,9 +327,7 @@ def test_synthesis_tool_names_run_python() -> None: AgentSimPredicateInventionApproach sim_learn = object.__new__(AgentSimLearningApproach) - sim_learn._do_synthesize_samplers = False invention = object.__new__(AgentSimPredicateInventionApproach) - invention._do_synthesize_samplers = False utils.reset_config({}) names = _required_names(sim_learn._get_synthesis_tool_names()) @@ -339,28 +336,6 @@ def test_synthesis_tool_names_run_python() -> None: assert names.count("run_python") == 1 assert "evaluate_predicate_quality" not in names # sim.predicates() - # On the solve side every arm with a simulator gets the same surface: - # the probe (trajectories in its namespace, sim.task for the task - # digest) plus the submission tool. - utils.reset_config({ - "env": "cover", - "approach": "agent_sim_predicate_invention", - "agent_planner_use_simulator": True, - }) - names = _required_names(invention._get_solve_tool_names()) - assert names.count("run_python") == 1 - assert "submit_plan" in names - - # Without a simulator there is nothing to probe or validate against. - utils.reset_config({ - "env": "cover", - "approach": "agent_sim_predicate_invention", - "agent_planner_use_simulator": False, - }) - names = _required_names(invention._get_solve_tool_names()) - assert "run_python" not in names - assert "submit_plan" not in names - def test_attached_run_python_replaces_the_static_instance() -> None: """A session that attaches its own ``run_python`` (synthesis) gets. diff --git a/tests/approaches/test_agent_model_based_approach.py b/tests/approaches/test_agent_model_based_approach.py deleted file mode 100644 index d0c7683aaf..0000000000 --- a/tests/approaches/test_agent_model_based_approach.py +++ /dev/null @@ -1,1919 +0,0 @@ -"""Tests for AgentModelBasedApproach -- parsing and refinement logic.""" -# pylint: disable=protected-access,import-outside-toplevel -import os -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest -from gym.spaces import Box - -from predicators import utils -from predicators.approaches.agent_model_based_approach import \ - AgentModelBasedApproach, _SketchStep -from predicators.structs import Action, GroundAtom, Object, \ - ParameterizedOption, Predicate, State, Task, Type - -_TEST_DATA_DIR = os.path.join(os.path.dirname(__file__), "test_data") - -# --------------------------------------------------------------------------- -# Shared fixtures -# --------------------------------------------------------------------------- - -_block_type = Type("block", ["x", "y", "held"]) -_robot_type = Type("robot", ["x", "y"]) - -_block0 = Object("block0", _block_type) -_block1 = Object("block1", _block_type) -_robot = Object("robot0", _robot_type) - -_Holding = Predicate("Holding", [_block_type], - lambda s, o: s.get(o[0], "held") > 0.5) -_On = Predicate("On", [_block_type, _block_type], - lambda s, o: abs(s.get(o[0], "x") - s.get(o[1], "x")) < 0.1) -_HandEmpty = Predicate("HandEmpty", [_robot_type], lambda s, o: True) - -_ALL_PREDICATES = {_Holding, _On, _HandEmpty} -_ALL_OBJECTS = [_block0, _block1, _robot] - - -def _noop_policy(_s, _m, _o, _p): - return Action(np.zeros(1, dtype=np.float32)) - - -def _always_true(_s, _m, _o, _p): - return True - - -def _always_false(_s, _m, _o, _p): - return False - - -_Pick = ParameterizedOption( - "Pick", - types=[_block_type], - params_space=Box(low=np.array([0.0], dtype=np.float32), - high=np.array([1.0], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_true, - terminal=_always_false, -) - -_Place = ParameterizedOption( - "Place", - types=[_block_type, _block_type], - params_space=Box(low=np.array([0.0, 0.0], dtype=np.float32), - high=np.array([1.0, 1.0], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_true, - terminal=_always_false, -) - -_Wait = ParameterizedOption( - "Wait", - types=[_robot_type], - params_space=Box(low=np.array([], dtype=np.float32), - high=np.array([], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_true, - terminal=_always_false, -) - -_ALL_OPTIONS = {_Pick, _Place, _Wait} - - -def _make_state(overrides=None): - """Create a simple state with default feature values.""" - data = { - _block0: np.array([0.1, 0.2, 0.0], dtype=np.float32), - _block1: np.array([0.5, 0.6, 0.0], dtype=np.float32), - _robot: np.array([0.0, 0.0], dtype=np.float32), - } - if overrides: - for obj, vals in overrides.items(): - data[obj] = np.array(vals, dtype=np.float32) - return State(data) - - -def _make_approach(): - """Create an AgentModelBasedApproach with mock config and option model.""" - state = _make_state() - goal = {GroundAtom(_On, [_block0, _block1])} - task = Task(state, goal) - - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "num_train_tasks": 1, - "num_test_tasks": 1, - "option_model_name": "oracle", - "seed": 42, - "agent_bilevel_max_samples_per_step": 10, - "agent_bilevel_check_subgoals": True, - }) - - mock_option_model = MagicMock() - approach = AgentModelBasedApproach( - initial_predicates=_ALL_PREDICATES, - initial_options=_ALL_OPTIONS, - types={_block_type, _robot_type}, - action_space=Box(low=-1, high=1, shape=(1, )), - train_tasks=[task], - option_model=mock_option_model, - ) - return approach, mock_option_model, task - - -# --------------------------------------------------------------------------- -# Tests: _parse_subgoal_annotations -# --------------------------------------------------------------------------- - - -class TestParseSubgoalAnnotations: - """Tests for plan text subgoal parsing.""" - - def test_basic_subgoals(self): - """Test basic subgoals.""" - approach, _, _ = _make_approach() - text = ("Pick(block0:block) -> {Holding(block0:block)}\n" - "Place(block0:block, block1:block) -> " - "{On(block0:block, block1:block)}\n") - result = approach._parse_subgoal_annotations(text, _ALL_PREDICATES, - _ALL_OBJECTS) - - assert len(result) == 2 - # First step: Holding(block0) - assert result[0] is not None - pos, neg = result[0] - assert GroundAtom(_Holding, [_block0]) in pos - assert len(neg) == 0 - # Second step: On(block0, block1) - assert result[1] is not None - pos2, neg2 = result[1] - assert GroundAtom(_On, [_block0, _block1]) in pos2 - assert len(neg2) == 0 - - def test_no_subgoals(self): - """Test no subgoals.""" - approach, _, _ = _make_approach() - text = ("Pick(block0:block)\n" - "Place(block0:block, block1:block)\n") - result = approach._parse_subgoal_annotations(text, _ALL_PREDICATES, - _ALL_OBJECTS) - - assert len(result) == 2 - assert result[0] is None - assert result[1] is None - - def test_mixed_subgoals(self): - """Some lines have subgoals, some don't.""" - approach, _, _ = _make_approach() - text = ("Pick(block0:block) -> {Holding(block0:block)}\n" - "Wait(robot0:robot)\n" - "Place(block0:block, block1:block) -> " - "{On(block0:block, block1:block)}\n") - result = approach._parse_subgoal_annotations(text, _ALL_PREDICATES, - _ALL_OBJECTS) - - assert len(result) == 3 - assert result[0] is not None - assert result[1] is None # Wait has no subgoal - assert result[2] is not None - - def test_multiple_atoms_in_subgoal(self): - """Test multiple atoms in subgoal.""" - approach, _, _ = _make_approach() - text = ( - "Place(block0:block, block1:block) " - "-> {On(block0:block, block1:block), HandEmpty(robot0:robot)}\n") - result = approach._parse_subgoal_annotations(text, _ALL_PREDICATES, - _ALL_OBJECTS) - - assert len(result) == 1 - assert result[0] is not None - pos, neg = result[0] - assert len(pos) == 2 - assert len(neg) == 0 - assert GroundAtom(_On, [_block0, _block1]) in pos - assert GroundAtom(_HandEmpty, [_robot]) in pos - - def test_unknown_predicate_skipped(self): - """Test unknown predicate skipped.""" - approach, _, _ = _make_approach() - text = "Pick(block0:block) -> {FakePred(block0:block)}\n" - result = approach._parse_subgoal_annotations(text, _ALL_PREDICATES, - _ALL_OBJECTS) - - assert len(result) == 1 - assert result[0] is None # FakePred unrecognized, no valid atoms - - def test_unknown_object_skipped(self): - """Test unknown object skipped.""" - approach, _, _ = _make_approach() - text = "Pick(block0:block) -> {Holding(block99:block)}\n" - result = approach._parse_subgoal_annotations(text, _ALL_PREDICATES, - _ALL_OBJECTS) - - assert len(result) == 1 - assert result[0] is None # block99 doesn't exist - - def test_arity_mismatch_skipped(self): - """Test arity mismatch skipped.""" - approach, _, _ = _make_approach() - # Holding expects 1 arg, giving 2 - text = "Pick(block0:block) -> {Holding(block0:block, block1:block)}\n" - result = approach._parse_subgoal_annotations(text, _ALL_PREDICATES, - _ALL_OBJECTS) - - assert len(result) == 1 - assert result[0] is None - - def test_typed_object_refs_in_subgoals(self): - """Agent outputs obj:type in subgoal atoms — should still parse.""" - approach, _, _ = _make_approach() - text = ("Pick(block0:block) -> {Holding(block0:block)}\n" - "Place(block0:block, block1:block) " - "-> {On(block0:block, block1:block)}\n") - result = approach._parse_subgoal_annotations(text, _ALL_PREDICATES, - _ALL_OBJECTS) - - assert len(result) == 2 - assert result[0] is not None - pos, _ = result[0] - assert GroundAtom(_Holding, [_block0]) in pos - assert result[1] is not None - pos2, _ = result[1] - assert GroundAtom(_On, [_block0, _block1]) in pos2 - - def test_numbered_prefix_subgoals(self): - """Agent numbers the lines (0:, 1:) — annotations must still align. - - Mirrors a real failure: the agent mirrored the numbered sketch - format shown in logs, embedding it between prose, and the - numbered prefix made every line parse as a non-option line so - the annotation list came back empty/misaligned. - """ - approach, _, _ = _make_approach() - text = ("Some analysis the agent wrote first.\n" - " 0: Pick(block0:block) -> {Holding(block0:block)}\n" - " 1: Place(block0:block, block1:block) " - "-> {On(block0:block, block1:block)}\n" - "Rationale: ...\n") - result = approach._parse_subgoal_annotations(text, _ALL_PREDICATES, - _ALL_OBJECTS) - - assert len(result) == 2 - assert result[0] is not None - pos, _ = result[0] - assert GroundAtom(_Holding, [_block0]) in pos - assert result[1] is not None - pos2, _ = result[1] - assert GroundAtom(_On, [_block0, _block1]) in pos2 - - def test_preamble_ignored(self): - """Non-option lines should be ignored.""" - approach, _, _ = _make_approach() - text = ("Here is my analysis:\n" - "I think we should pick block0 first.\n" - "\n" - "Pick(block0:block) -> {Holding(block0:block)}\n") - result = approach._parse_subgoal_annotations(text, _ALL_PREDICATES, - _ALL_OBJECTS) - - assert len(result) == 1 - assert result[0] is not None - - def test_whitespace_in_atoms(self): - """Spaces around commas in atom arguments.""" - approach, _, _ = _make_approach() - text = ("Place(block0:block, block1:block) -> " - "{ On( block0:block , block1:block ) }\n") - result = approach._parse_subgoal_annotations(text, _ALL_PREDICATES, - _ALL_OBJECTS) - - assert len(result) == 1 - assert result[0] is not None - pos, _ = result[0] - assert GroundAtom(_On, [_block0, _block1]) in pos - - def test_not_atoms_in_subgoals(self): - """Test NOT prefix for negative target atoms.""" - approach, _, _ = _make_approach() - text = ( - "Wait(robot0:robot) -> " - "{Holding(block0:block), NOT On(block0:block, block1:block)}\n") - result = approach._parse_subgoal_annotations(text, _ALL_PREDICATES, - _ALL_OBJECTS) - - assert len(result) == 1 - assert result[0] is not None - pos, neg = result[0] - assert GroundAtom(_Holding, [_block0]) in pos - assert GroundAtom(_On, [_block0, _block1]) in neg - - -# --------------------------------------------------------------------------- -# Tests: check_wait_target_atoms -# --------------------------------------------------------------------------- - - -class TestCheckWaitTargetAtoms: - """Tests that Wait terminates on target atoms, not noisy changes.""" - - def test_no_targets_returns_none(self): - """No targets in memory -> returns None (fall back to any-change).""" - opt = _Wait.ground([_robot], np.array([], dtype=np.float32)) - # No targets in memory - state = _make_state({_block0: [0.0, 0.0, 0.0]}) - abstract_fn = lambda s: utils.abstract(s, _ALL_PREDICATES) - result = utils.check_wait_target_atoms(opt, state, abstract_fn) - assert result is None - - def test_positive_target_met(self): - """Wait terminates when positive target atom holds.""" - opt = _Wait.ground([_robot], np.array([], dtype=np.float32)) - target_atom = GroundAtom(_Holding, [_block0]) - opt.memory["wait_target_atoms"] = {target_atom} - - # State where Holding(block0) is true (held > 0.5) - state_held = _make_state({_block0: [0.0, 0.0, 1.0]}) - abstract_fn = lambda s: utils.abstract(s, _ALL_PREDICATES) - assert utils.check_wait_target_atoms(opt, state_held, abstract_fn) \ - is True - - def test_positive_target_not_met(self): - """Wait does NOT terminate when target atom doesn't hold yet.""" - opt = _Wait.ground([_robot], np.array([], dtype=np.float32)) - target_atom = GroundAtom(_Holding, [_block0]) - opt.memory["wait_target_atoms"] = {target_atom} - - # State where Holding(block0) is false (held <= 0.5) - state_not_held = _make_state({_block0: [0.0, 0.0, 0.0]}) - abstract_fn = lambda s: utils.abstract(s, _ALL_PREDICATES) - assert utils.check_wait_target_atoms(opt, state_not_held, - abstract_fn) is False - - def test_noisy_atom_change_ignored_with_targets(self): - """Wait ignores noisy atom changes when specific targets are set. - - This is the key test: if the Wait is parameterized with a target - atom (e.g. Holding(block0)), it should NOT terminate when a - different atom changes (e.g. On(block0, block1)). - """ - opt = _Wait.ground([_robot], np.array([], dtype=np.float32)) - # Only waiting for Holding(block0) - target_atom = GroundAtom(_Holding, [_block0]) - opt.memory["wait_target_atoms"] = {target_atom} - - # State where On(block0, block1) is true (noisy change) but - # Holding(block0) is still false - state_noisy = _make_state({ - _block0: [0.5, 0.0, 0.0], - _block1: [0.5, 0.0, 0.0] - }) - abstract_fn = lambda s: utils.abstract(s, _ALL_PREDICATES) - atoms = abstract_fn(state_noisy) - # On is true (positions are close), but Holding is false - assert GroundAtom(_On, [_block0, _block1]) in atoms - assert GroundAtom(_Holding, [_block0]) not in atoms - - # Wait should NOT terminate (target not met, despite On changing) - assert utils.check_wait_target_atoms(opt, state_noisy, - abstract_fn) is False - - def test_negative_target_met(self): - """Wait terminates when negative target atom is false.""" - opt = _Wait.ground([_robot], np.array([], dtype=np.float32)) - neg_atom = GroundAtom(_On, [_block0, _block1]) - opt.memory["wait_target_neg_atoms"] = {neg_atom} - - # State where On(block0, block1) is false (positions far apart) - state = _make_state({ - _block0: [0.0, 0.0, 0.0], - _block1: [5.0, 0.0, 0.0] - }) - abstract_fn = lambda s: utils.abstract(s, _ALL_PREDICATES) - assert utils.check_wait_target_atoms(opt, state, abstract_fn) is True - - def test_negative_target_not_met(self): - """Wait does NOT terminate when negative target atom is still true.""" - opt = _Wait.ground([_robot], np.array([], dtype=np.float32)) - neg_atom = GroundAtom(_On, [_block0, _block1]) - opt.memory["wait_target_neg_atoms"] = {neg_atom} - - # State where On(block0, block1) is true (positions close) - state = _make_state({ - _block0: [0.5, 0.0, 0.0], - _block1: [0.5, 0.0, 0.0] - }) - abstract_fn = lambda s: utils.abstract(s, _ALL_PREDICATES) - assert utils.check_wait_target_atoms(opt, state, abstract_fn) is False - - def test_mixed_positive_and_negative_targets(self): - """Both positive and negative targets must be satisfied.""" - opt = _Wait.ground([_robot], np.array([], dtype=np.float32)) - opt.memory["wait_target_atoms"] = {GroundAtom(_Holding, [_block0])} - opt.memory["wait_target_neg_atoms"] = { - GroundAtom(_On, [_block0, _block1]) - } - - abstract_fn = lambda s: utils.abstract(s, _ALL_PREDICATES) - - # Only positive met (Holding true, On still true) - state1 = _make_state({ - _block0: [0.5, 0.0, 1.0], - _block1: [0.5, 0.0, 0.0] - }) - assert utils.check_wait_target_atoms(opt, state1, abstract_fn) is False - - # Only negative met (On false, Holding false) - state2 = _make_state({ - _block0: [0.0, 0.0, 0.0], - _block1: [5.0, 0.0, 0.0] - }) - assert utils.check_wait_target_atoms(opt, state2, abstract_fn) is False - - # Both met (Holding true, On false) - state3 = _make_state({ - _block0: [0.0, 0.0, 1.0], - _block1: [5.0, 0.0, 0.0] - }) - assert utils.check_wait_target_atoms(opt, state3, abstract_fn) is True - - -# --------------------------------------------------------------------------- -# Tests: parse_wait_target_annotations and strip_wait_annotations -# --------------------------------------------------------------------------- - - -class TestWaitTargetParsing: - """Tests for parse_wait_target_annotations and strip_wait_annotations.""" - - def test_parse_positive_target(self): - """Parse a positive target atom.""" - line = "Wait(robot0:robot) -> {Holding(block0:block)}" - pos, neg = utils.parse_wait_target_annotations(line, _ALL_PREDICATES, - _ALL_OBJECTS) - assert GroundAtom(_Holding, [_block0]) in pos - assert len(neg) == 0 - - def test_parse_negative_target(self): - """Parse a NOT-prefixed target atom.""" - line = "Wait(robot0:robot) -> {NOT On(block0:block, block1:block)}" - pos, neg = utils.parse_wait_target_annotations(line, _ALL_PREDICATES, - _ALL_OBJECTS) - assert len(pos) == 0 - assert GroundAtom(_On, [_block0, _block1]) in neg - - def test_parse_mixed_targets(self): - """Parse both positive and negative target atoms.""" - line = ("Wait(robot0:robot) -> " - "{Holding(block0:block), NOT On(block0:block, block1:block)}") - pos, neg = utils.parse_wait_target_annotations(line, _ALL_PREDICATES, - _ALL_OBJECTS) - assert GroundAtom(_Holding, [_block0]) in pos - assert GroundAtom(_On, [_block0, _block1]) in neg - - def test_parse_no_annotation(self): - """Line without -> returns empty sets.""" - line = "Wait(robot0:robot)[]" - pos, neg = utils.parse_wait_target_annotations(line, _ALL_PREDICATES, - _ALL_OBJECTS) - assert len(pos) == 0 - assert len(neg) == 0 - - def test_strip_annotations(self): - """strip_wait_annotations removes -> {...} suffixes.""" - text = ("Pick(block0:block)[0.5]\n" - "Wait(robot0:robot)[] -> {Holding(block0:block)}\n" - "Place(block0:block, block1:block)[0.1, 0.2]\n") - stripped = utils.strip_wait_annotations(text) - assert "-> {" not in stripped - assert "Pick(block0:block)[0.5]" in stripped - assert "Wait(robot0:robot)[]" in stripped - assert "Place(block0:block, block1:block)[0.1, 0.2]" in stripped - - -# --------------------------------------------------------------------------- -# Tests: _refine_sketch -# --------------------------------------------------------------------------- - - -class TestRefineSketch: - """Tests for backtracking refinement search.""" - - def test_empty_sketch(self): - """Test empty sketch.""" - approach, _, task = _make_approach() - plan, success = approach._refine_sketch(task, [], timeout=5.0) - assert plan == [] - assert success is False - - def test_single_step_no_params(self): - """Option with empty params_space — should succeed in 1 try.""" - approach, mock_om, task = _make_approach() - - # Option model: Wait always succeeds, goal holds after - goal_state = _make_state({_block0: [0.5, 0.6, 0.0]}) - mock_om.get_next_state_and_num_actions.return_value = (goal_state, 5) - - sketch = [ - _SketchStep(option=_Wait, objects=[_robot], subgoal_atoms=None) - ] - plan, success = approach._refine_sketch(task, sketch, timeout=5.0) - - assert success is True - assert len(plan) == 1 - assert plan[0].name == "Wait" - - def test_single_step_with_params_success(self): - """Option with params — should find working params via sampling.""" - approach, mock_om, task = _make_approach() - - goal_state = _make_state({_block0: [0.5, 0.6, 0.0]}) - mock_om.get_next_state_and_num_actions.return_value = (goal_state, 3) - - sketch = [ - _SketchStep(option=_Pick, objects=[_block0], subgoal_atoms=None) - ] - plan, success = approach._refine_sketch(task, sketch, timeout=5.0) - - assert success is True - assert len(plan) == 1 - - def test_subgoal_check_pass(self): - """Subgoal atoms hold after execution.""" - approach, mock_om, task = _make_approach() - - # After Pick, Holding(block0) should hold — set held=1 - held_state = _make_state({_block0: [0.1, 0.2, 1.0]}) - # After Place, On(block0, block1) — set x close - goal_state = _make_state({_block0: [0.5, 0.6, 0.0]}) - - mock_om.get_next_state_and_num_actions.side_effect = [ - (held_state, 3), - (goal_state, 3), - ] - - sketch = [ - _SketchStep(option=_Pick, - objects=[_block0], - subgoal_atoms={GroundAtom(_Holding, [_block0])}), - _SketchStep(option=_Place, - objects=[_block0, _block1], - subgoal_atoms={GroundAtom(_On, [_block0, _block1])}), - ] - plan, success = approach._refine_sketch(task, sketch, timeout=5.0) - - assert success is True - assert len(plan) == 2 - - def test_subgoal_check_fail_triggers_resample(self): - """Subgoal atoms don't hold — should resample params.""" - approach, mock_om, task = _make_approach() - - # Holding never holds (held=0) — subgoal always fails - bad_state = _make_state({_block0: [0.1, 0.2, 0.0]}) - mock_om.get_next_state_and_num_actions.return_value = (bad_state, 3) - - sketch = [ - _SketchStep(option=_Pick, - objects=[_block0], - subgoal_atoms={GroundAtom(_Holding, [_block0])}), - ] - _plan, success = approach._refine_sketch(task, sketch, timeout=5.0) - - # Should exhaust all samples and fail - assert success is False - # Option model called max_samples times (10) - assert mock_om.get_next_state_and_num_actions.call_count == 10 - - def test_backtracking_across_steps(self): - """Step 2 fails, causing step 1 to be re-sampled.""" - approach, mock_om, task = _make_approach() - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "num_train_tasks": 1, - "num_test_tasks": 1, - "seed": 42, - "agent_bilevel_max_samples_per_step": 3, - "agent_bilevel_check_subgoals": False, - }) - - call_count = 0 - goal_state = _make_state({_block0: [0.5, 0.6, 0.0]}) - noop_state = _make_state() - - def side_effect(_state, option): - nonlocal call_count - call_count += 1 - if option.name == "Pick": - return (noop_state, 3) # Pick always succeeds - # Place: succeed only on the last attempt - if call_count >= 8: - return (goal_state, 3) - return (noop_state, 0) # fail (noop) - - mock_om.get_next_state_and_num_actions.side_effect = side_effect - - sketch = [ - _SketchStep(option=_Pick, objects=[_block0], subgoal_atoms=None), - _SketchStep(option=_Place, - objects=[_block0, _block1], - subgoal_atoms=None), - ] - plan, success = approach._refine_sketch(task, sketch, timeout=10.0) - - # Should have backtracked and eventually succeeded - assert success is True - assert len(plan) == 2 - assert call_count >= 4 # at least one backtrack cycle - - def test_not_initiable_triggers_resample(self): - """Option not initiable in current state — resample.""" - approach, mock_om, task = _make_approach() - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "num_train_tasks": 1, - "num_test_tasks": 1, - "seed": 42, - "agent_bilevel_max_samples_per_step": 3, - }) - - # Create an option that is never initiable - not_initiable = ParameterizedOption( - "Pick", - types=[_block_type], - params_space=Box(low=np.array([0.0], dtype=np.float32), - high=np.array([1.0], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_false, - terminal=_always_false, - ) - - sketch = [ - _SketchStep(option=not_initiable, - objects=[_block0], - subgoal_atoms=None) - ] - _plan, success = approach._refine_sketch(task, sketch, timeout=5.0) - - assert success is False - # Option model never called since initiable is always False - mock_om.get_next_state_and_num_actions.assert_not_called() - - def test_goal_check_on_final_step(self): - """Final step must satisfy the task goal even without subgoals.""" - approach, mock_om, task = _make_approach() - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "num_train_tasks": 1, - "num_test_tasks": 1, - "seed": 42, - "agent_bilevel_max_samples_per_step": 5, - "agent_bilevel_check_subgoals": False, - }) - - # State that doesn't satisfy goal On(block0, block1) - bad_state = _make_state({_block0: [0.9, 0.2, 0.0]}) - mock_om.get_next_state_and_num_actions.return_value = (bad_state, 3) - - sketch = [ - _SketchStep(option=_Pick, objects=[_block0], subgoal_atoms=None) - ] - _plan, success = approach._refine_sketch(task, sketch, timeout=5.0) - - # Goal never holds → exhausts samples - assert success is False - - -# --------------------------------------------------------------------------- -# Tests: _query_agent_for_plan_sketch (with mocked agent) -# --------------------------------------------------------------------------- - - -class TestQueryAgentForPlanSketch: - """Tests for end-to-end sketch extraction from mock agent responses.""" - - def _mock_responses(self, plan_text): - """Build mock agent response list containing plan_text.""" - return [ - { - "type": "assistant", - "content": [{ - "type": "text", - "text": plan_text - }], - }, - ] - - def test_basic_sketch_extraction(self): - """Test basic sketch extraction.""" - approach, _, task = _make_approach() - - plan_text = ("Pick(block0:block) -> {Holding(block0:block)}\n" - "Place(block0:block, block1:block) -> " - "{On(block0:block, block1:block)}\n") - - with patch.object(approach, - '_query_agent_sync', - return_value=self._mock_responses(plan_text)): - sketch = approach._query_agent_for_plan_sketch(task) - - assert len(sketch) == 2 - assert sketch[0].option.name == "Pick" - assert list(sketch[0].objects) == [_block0] - assert sketch[0].subgoal_atoms is not None - assert GroundAtom(_Holding, [_block0]) in sketch[0].subgoal_atoms - - assert sketch[1].option.name == "Place" - assert list(sketch[1].objects) == [_block0, _block1] - assert sketch[1].subgoal_atoms is not None - - def test_sketch_without_subgoals(self): - """Test sketch without subgoals.""" - approach, _, task = _make_approach() - - plan_text = ("Pick(block0:block)\n" - "Place(block0:block, block1:block)\n") - - with patch.object(approach, - '_query_agent_sync', - return_value=self._mock_responses(plan_text)): - sketch = approach._query_agent_for_plan_sketch(task) - - assert len(sketch) == 2 - assert sketch[0].subgoal_atoms is None - assert sketch[1].subgoal_atoms is None - - def test_sketch_with_code_fences(self): - """Test sketch with code fences.""" - approach, _, task = _make_approach() - - plan_text = ("Here is the plan:\n" - "```\n" - "Pick(block0:block) -> {Holding(block0:block)}\n" - "Place(block0:block, block1:block)\n" - "```\n") - - with patch.object(approach, - '_query_agent_sync', - return_value=self._mock_responses(plan_text)): - sketch = approach._query_agent_for_plan_sketch(task) - - assert len(sketch) == 2 - - def test_sketch_with_preamble(self): - """Agent includes analysis text before the plan.""" - approach, _, task = _make_approach() - - plan_text = ( - "After inspecting the environment, I found block0 and block1.\n" - "The goal is to place block0 on block1.\n" - "\n" - "Pick(block0:block)\n" - "Place(block0:block, block1:block)\n") - - with patch.object(approach, - '_query_agent_sync', - return_value=self._mock_responses(plan_text)): - sketch = approach._query_agent_for_plan_sketch(task) - - assert len(sketch) == 2 - - def test_sketch_with_wait(self): - """Test sketch with wait.""" - approach, _, task = _make_approach() - - plan_text = ("Pick(block0:block) -> {Holding(block0:block)}\n" - "Wait(robot0:robot)\n" - "Place(block0:block, block1:block) -> " - "{On(block0:block, block1:block)}\n") - - with patch.object(approach, - '_query_agent_sync', - return_value=self._mock_responses(plan_text)): - sketch = approach._query_agent_for_plan_sketch(task) - - assert len(sketch) == 3 - assert sketch[0].option.name == "Pick" - assert sketch[1].option.name == "Wait" - assert sketch[1].subgoal_atoms is None - assert sketch[2].option.name == "Place" - - def test_empty_response_raises(self): - """Agent returns no text → ApproachFailure.""" - from predicators.approaches import ApproachFailure - approach, _, task = _make_approach() - - with patch.object(approach, - '_query_agent_sync', - return_value=[{ - "type": "result", - "content": [] - }]): - with pytest.raises(ApproachFailure, match="empty plan text"): - approach._query_agent_for_plan_sketch(task) - - def test_no_valid_options_raises(self): - """Agent returns text with no valid option names → ApproachFailure.""" - from predicators.approaches import ApproachFailure - approach, _, task = _make_approach() - - plan_text = "I don't know what to do.\nSorry!\n" - - with patch.object(approach, - '_query_agent_sync', - return_value=self._mock_responses(plan_text)): - with pytest.raises(ApproachFailure, match="Parsed empty"): - approach._query_agent_for_plan_sketch(task) - - def test_sketch_from_file(self): - """Load sketch from a saved text file via CFG option.""" - approach, _, task = _make_approach() - sketch_path = os.path.join(_TEST_DATA_DIR, "simple_plan_sketch.txt") - - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "num_train_tasks": 1, - "num_test_tasks": 1, - "seed": 42, - "agent_bilevel_plan_sketch_file": sketch_path, - }) - - sketch = approach._query_agent_for_plan_sketch(task) - - assert len(sketch) == 2 - assert sketch[0].option.name == "Pick" - assert list(sketch[0].objects) == [_block0] - assert sketch[0].subgoal_atoms is not None - assert GroundAtom(_Holding, [_block0]) in sketch[0].subgoal_atoms - assert sketch[1].option.name == "Place" - assert list(sketch[1].objects) == [_block0, _block1] - assert sketch[1].subgoal_atoms is not None - assert GroundAtom(_On, [_block0, _block1]) in sketch[1].subgoal_atoms - - -# --------------------------------------------------------------------------- -# Tests: _sample_params -# --------------------------------------------------------------------------- - - -class TestValidatePlanForward: - """Tests for ``plan_execution.validate_plan_forward``. - - Covers the test-time forward validator that's the entire reason the - synthesis tool can catch refinement-passes/validation-fails - regressions. - """ - - def _grounded(self, option, objects, params=None): - if params is None: - params = np.zeros(option.params_space.shape[0], dtype=np.float32) - return option.ground(list(objects), np.asarray(params, - dtype=np.float32)) - - def test_goal_reached_returns_success(self): - """Plan that reaches the goal — validator passes, no diagnosis.""" - from predicators.agent_sdk import plan_execution - _, mock_om, task = _make_approach() - # Final post-state satisfies the goal (On(block0, block1)). - goal_state = _make_state({_block0: [0.55, 0.6, 0.0]}) - mock_om.get_next_state_and_num_actions.return_value = (goal_state, 3) - - plan = [self._grounded(_Pick, [_block0], [0.5])] - ok, reason = plan_execution.validate_plan_forward( - task, plan, mock_om, predicates=_ALL_PREDICATES) - assert ok is True - assert reason == "" - - def test_goal_not_reached_diagnosis_names_missing_atoms(self): - """Plan terminates but goal isn't satisfied — diagnosis names the - missing atom set, not a generic 'validation failed'.""" - from predicators.agent_sdk import plan_execution - _, mock_om, task = _make_approach() - # Post-state doesn't satisfy On(block0, block1). - bad_state = _make_state({_block0: [0.1, 0.2, 0.0]}) - mock_om.get_next_state_and_num_actions.return_value = (bad_state, 3) - - plan = [self._grounded(_Pick, [_block0], [0.5])] - ok, reason = plan_execution.validate_plan_forward( - task, plan, mock_om, predicates=_ALL_PREDICATES) - assert ok is False - assert "goal not reached" in reason - assert "On(block0:block, block1:block)" in reason - - def test_subgoal_divergence_logged_when_sketch_provided(self, caplog): - """When the sketch is passed in, per-step subgoal divergence is logged - with the missing atom — this is the diagnostic the synthesis agent - needs to see *which* step's predicate is spurious.""" - import logging as _logging - - from predicators.agent_sdk import plan_execution - _, mock_om, task = _make_approach() - # Post-state never establishes Holding(block0). Goal is also - # missing — but the subgoal log should fire first. - bad_state = _make_state({_block0: [0.1, 0.2, 0.0]}) - mock_om.get_next_state_and_num_actions.return_value = (bad_state, 3) - - plan = [self._grounded(_Pick, [_block0], [0.5])] - sketch = [ - _SketchStep(option=_Pick, - objects=[_block0], - subgoal_atoms={GroundAtom(_Holding, [_block0])}) - ] - with caplog.at_level(_logging.INFO): - ok, _ = plan_execution.validate_plan_forward( - task, - plan, - mock_om, - predicates=_ALL_PREDICATES, - sketch=sketch, - run_id="test_run", - ) - assert ok is False - # Subgoal divergence log mentions the missing atom and the step. - assert any("subgoal divergence at step 0" in r.message - and "Holding(block0:block)" in r.message - for r in caplog.records) - - def test_option_failure_diagnosis_names_step(self): - """When the option model returns 0 actions (option execution failed), - the diagnosis identifies the failing step and surfaces the option - model's last_execution_failure.""" - from predicators.agent_sdk import plan_execution - _, mock_om, task = _make_approach() - # Simulate option failure: 0 actions, with a diagnostic message - # recorded on the option model. - mock_om.get_next_state_and_num_actions.return_value = (_make_state(), - 0) - mock_om.last_execution_failure = "IK timed out at waypoint 3" - - plan = [self._grounded(_Pick, [_block0], [0.5])] - ok, reason = plan_execution.validate_plan_forward( - task, plan, mock_om, predicates=_ALL_PREDICATES) - assert ok is False - assert "option execution failed at step 0" in reason - assert "Pick(block0)" in reason - assert "IK timed out at waypoint 3" in reason - - def test_empty_plan_with_goal_already_satisfied(self): - """Empty plan + init satisfies goal → success.""" - from predicators.agent_sdk import plan_execution - - # Goal trivially holds when block0 is already on block1. - init = _make_state({_block0: [0.55, 0.6, 0.0]}) - task = Task(init, {GroundAtom(_On, [_block0, _block1])}) - mock_om = MagicMock() - ok, reason = plan_execution.validate_plan_forward( - task, [], mock_om, predicates=_ALL_PREDICATES) - assert ok is True - assert reason == "" - - def test_empty_plan_with_unmet_goal(self): - """Empty plan + init does NOT satisfy goal → failure with explanatory - diagnosis.""" - from predicators.agent_sdk import plan_execution - _, _, task = _make_approach() # init does not satisfy goal - mock_om = MagicMock() - ok, reason = plan_execution.validate_plan_forward( - task, [], mock_om, predicates=_ALL_PREDICATES) - assert ok is False - assert "init state does not satisfy goal" in reason - - def test_sketch_length_mismatch_ignored_gracefully(self): - """Mismatched sketch length — validator should warn and fall back to - goal-only checking rather than crash.""" - from predicators.agent_sdk import plan_execution - _, mock_om, task = _make_approach() - goal_state = _make_state({_block0: [0.55, 0.6, 0.0]}) - mock_om.get_next_state_and_num_actions.return_value = (goal_state, 3) - - plan = [self._grounded(_Pick, [_block0], [0.5])] - # Sketch length 2, plan length 1. - sketch = [ - _SketchStep(option=_Pick, objects=[_block0], subgoal_atoms=None), - _SketchStep(option=_Pick, objects=[_block0], subgoal_atoms=None), - ] - ok, _ = plan_execution.validate_plan_forward( - task, - plan, - mock_om, - predicates=_ALL_PREDICATES, - sketch=sketch, - ) - # Validation still runs to completion against the goal. - assert ok is True - - -class TestSampleParams: - """TestSampleParams class.""" - - def test_empty_params_space(self): - """Test empty params space.""" - approach, _, _ = _make_approach() - rng = np.random.default_rng(0) - params = approach._sample_params(_Wait, _make_state(), rng) - assert params.shape == (0, ) - assert params.dtype == np.float32 - - def test_params_within_bounds(self): - """Test params within bounds.""" - approach, _, _ = _make_approach() - rng = np.random.default_rng(0) - for _ in range(100): - params = approach._sample_params(_Place, _make_state(), rng) - assert params.shape == (2, ) - assert np.all(params >= 0.0) - assert np.all(params <= 1.0) - assert params.dtype == np.float32 - - -# --------------------------------------------------------------------------- -# Tests: class metadata -# --------------------------------------------------------------------------- - - -def test_get_name(): - """Test get name.""" - assert AgentModelBasedApproach.get_name() == "agent_model_based" - # The pre-rename CLI name still resolves to this approach. - # pylint: disable=import-outside-toplevel - from predicators.approaches import _get_approach_cls_from_name - assert _get_approach_cls_from_name( - "agent_bilevel") is AgentModelBasedApproach - - -# --------------------------------------------------------------------------- -# Tests: closed-loop execution replanning (subgoal_annotations monitor + -# _maybe_replan_from_divergence / _replan_suffix) -# --------------------------------------------------------------------------- - -_PickDone = ParameterizedOption( - "Pick", - types=[_block_type], - params_space=Box(low=np.array([0.0], dtype=np.float32), - high=np.array([1.0], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_true, - terminal=_always_true, -) - -_PlaceDone = ParameterizedOption( - "Place", - types=[_block_type, _block_type], - params_space=Box(low=np.array([0.0, 0.0], dtype=np.float32), - high=np.array([1.0, 1.0], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_true, - terminal=_always_true, -) - - -def _make_two_step_plan(first_subgoals): - """Plan [Pick, Place] whose first step is annotated with first_subgoals.""" - plan = [ - _PickDone.ground([_block0], np.array([0.5], dtype=np.float32)), - _PlaceDone.ground([_block0, _block1], - np.array([0.5, 0.5], dtype=np.float32)), - ] - sketch = [ - _SketchStep(_PickDone, [_block0], first_subgoals), - _SketchStep(_PlaceDone, [_block0, _block1], None), - ] - return plan, sketch - - -def _enable_replanning(approach, budget): - """Turn on closed-loop execution and start a fresh episode.""" - utils.update_config({ - "agent_bilevel_max_execution_replans": budget, - "execution_monitor": "subgoal_annotations", - }) - approach.reset_for_new_episode() - - -def _make_monitor(approach): - """Create the monitor and sync it with the approach, CogMan-style.""" - from predicators.execution_monitoring import create_execution_monitor - monitor = create_execution_monitor("subgoal_annotations") - monitor.update_approach_info(approach.get_execution_monitoring_info()) - return monitor - - -def _sync(monitor, approach): - """Mimic CogMan pushing fresh approach info to the monitor.""" - monitor.update_approach_info(approach.get_execution_monitoring_info()) - - -class TestExecutionReplanning: - """Tests for closed-loop execution through the cogman monitor flow.""" - - def test_open_loop_when_disabled(self): - """With the flag at 0 (default), no monitoring info is exported and - divergence is never flagged.""" - approach, _, _ = _make_approach() - holding = {GroundAtom(_Holding, [_block0])} - plan, sketch = _make_two_step_plan(holding) - policy = approach._plan_to_policy(plan, sketch=sketch) - assert not approach.get_execution_monitoring_info() - state = _make_state() # block0 not held: subgoal would fail - monitor = _make_monitor(approach) - assert not monitor.step(state) - policy(state) # starts Pick - policy(state) # Pick terminal -> starts Place without any check - - def test_monitor_silent_when_subgoals_hold(self): - """Subgoals satisfied at the boundary: no replan is suggested.""" - approach, _, _ = _make_approach() - _enable_replanning(approach, 2) - holding = {GroundAtom(_Holding, [_block0])} - plan, sketch = _make_two_step_plan(holding) - policy = approach._plan_to_policy(plan, sketch=sketch) - state = _make_state({_block0: [0.1, 0.2, 1.0]}) # held: subgoal ok - monitor = _make_monitor(approach) - # Before any option is initiated (e.g. right after a replan, - # cogman asserts the monitor does not immediately re-fire). - assert not monitor.step(state) - policy(state) # starts Pick - _sync(monitor, approach) - assert not monitor.step(state) # boundary, but annotation holds - policy(state) # advances to Place - - def test_monitor_silent_mid_option(self): - """A failing annotation is only checked at the option boundary.""" - approach, _, _ = _make_approach() - _enable_replanning(approach, 2) - holding = {GroundAtom(_Holding, [_block0])} - # _Pick never terminates, so execution stays mid-option. - plan = [_Pick.ground([_block0], np.array([0.5], dtype=np.float32))] - sketch = [_SketchStep(_Pick, [_block0], holding)] - policy = approach._plan_to_policy(plan, sketch=sketch) - state = _make_state() # block0 not held: subgoal fails - policy(state) - monitor = _make_monitor(approach) - assert not monitor.step(state) - - def test_monitor_detects_divergence_at_boundary(self): - """An unsatisfied annotation at the boundary suggests a replan.""" - approach, _, _ = _make_approach() - _enable_replanning(approach, 2) - holding = {GroundAtom(_Holding, [_block0])} - plan, sketch = _make_two_step_plan(holding) - policy = approach._plan_to_policy(plan, sketch=sketch) - state = _make_state() # block0 not held: subgoal diverges - policy(state) # starts Pick (terminal at every state) - monitor = _make_monitor(approach) - assert monitor.step(state) - - def test_suffix_replan_preferred_on_divergence(self): - """The monitor-triggered re-solve resumes via the suffix path; no agent - re-query.""" - approach, _, task = _make_approach() - _enable_replanning(approach, 2) - holding = {GroundAtom(_Holding, [_block0])} - plan, sketch = _make_two_step_plan(holding) - policy = approach._plan_to_policy(plan, sketch=sketch) - state = _make_state() - policy(state) - monitor = _make_monitor(approach) - assert monitor.step(state) - - # CogMan now re-invokes solve() on the current state. - def sentinel_policy(s): - del s # unused - return Action(np.full(1, 0.25, dtype=np.float32)) - - approach._replan_suffix = MagicMock(return_value=sentinel_policy) - approach._query_agent_for_plan_sketch = MagicMock() - new_policy = approach._solve(Task(state, task.goal), timeout=10) - assert new_policy is sentinel_policy - approach._query_agent_for_plan_sketch.assert_not_called() - approach._replan_suffix.assert_called_once() - args = approach._replan_suffix.call_args.args - assert args[0] is state # replans from the real current state - assert args[3] == 0 # the failed step is the annotated first step - - def test_openloop_resume_when_no_suffix_validates(self): - """Suffix path exhausted: the remaining plan resumes open-loop. - - An annotation is the agent's prediction, not proof the goal is - out of reach, so by default the episode keeps executing and the - goal check decides. A fresh sketch query would re-open the agent - turn budget the attempt already spent, so it stays opt-in - (agent_bilevel_replan_agent_fallback). - """ - approach, _, task = _make_approach() - _enable_replanning(approach, 2) - holding = {GroundAtom(_Holding, [_block0])} - plan, sketch = _make_two_step_plan(holding) - policy = approach._plan_to_policy(plan, sketch=sketch) - state = _make_state() - policy(state) - approach._replan_suffix = MagicMock(return_value=None) - approach._query_agent_for_plan_sketch = MagicMock() - new_policy = approach._solve(Task(state, task.goal), timeout=10) - approach._replan_suffix.assert_called_once() - approach._query_agent_for_plan_sketch.assert_not_called() - # The resumed policy executes the remaining step (Place), and - # monitoring re-arms over exactly that suffix. - new_policy(state) - status = approach.get_execution_monitoring_info()[0] - assert status.steps_initiated == 1 - assert status.current_option.name == "Place" - - def test_openloop_resume_at_last_step_ends_plan(self): - """Divergence at the final step leaves nothing to resume: the returned - policy ends through the normal plan-exhausted path (so a goal-reached - terminator still gets its chance), not a divergence abort.""" - from predicators.approaches import ApproachFailure - approach, _, task = _make_approach() - _enable_replanning(approach, 2) - holding = {GroundAtom(_Holding, [_block0])} - plan, _ = _make_two_step_plan(holding) - # Annotate the LAST step instead of the first. - sketch = [ - _SketchStep(_PickDone, [_block0], None), - _SketchStep(_PlaceDone, [_block0, _block1], holding), - ] - policy = approach._plan_to_policy(plan, sketch=sketch) - state = _make_state() # block0 not held: Place's subgoal fails - policy(state) # starts Pick - policy(state) # Pick terminal -> starts Place - monitor = _make_monitor(approach) - assert monitor.step(state) - approach._replan_suffix = MagicMock(return_value=None) - new_policy = approach._solve(Task(state, task.goal), timeout=10) - with pytest.raises(ApproachFailure, match="exhausted"): - new_policy(state) - - def test_full_resolve_when_no_suffix_validates_with_fallback(self): - """With agent_bilevel_replan_agent_fallback, a failed suffix replan - falls through to a fresh agent sketch.""" - from predicators.approaches import ApproachFailure - approach, _, task = _make_approach() - _enable_replanning(approach, 2) - utils.update_config({"agent_bilevel_replan_agent_fallback": True}) - holding = {GroundAtom(_Holding, [_block0])} - plan, sketch = _make_two_step_plan(holding) - policy = approach._plan_to_policy(plan, sketch=sketch) - state = _make_state() - policy(state) - approach._replan_suffix = MagicMock(return_value=None) - # Reaching the fresh-sketch body raises its distinctive failure - - # proof we fell through to a fresh agent query. - sketch_query = MagicMock(side_effect=ApproachFailure("no sketch")) - approach._query_agent_for_plan_sketch = sketch_query - with patch.object(approach, '_nudge_final_submission', - MagicMock(return_value=None)): - with pytest.raises(ApproachFailure, match="Bilevel solve failed"): - approach._solve(Task(state, task.goal), timeout=10) - sketch_query.assert_called_once() - approach._replan_suffix.assert_called_once() - - def test_budget_shared_across_chained_replans(self): - """Chained replans share one per-episode budget; once it is exhausted, - a divergence resumes the remaining plan open-loop without paying for - further refinement.""" - approach, _, task = _make_approach() - _enable_replanning(approach, 1) - holding = {GroundAtom(_Holding, [_block0])} - plan, sketch = _make_two_step_plan(holding) - - def _suffix_replan(s, tsk, steps, k, t): - del s, tsk, steps, k, t # unused - new_plan, new_sketch = _make_two_step_plan(holding) - return approach._plan_to_policy(new_plan, sketch=new_sketch) - - approach._replan_suffix = MagicMock(side_effect=_suffix_replan) - approach._query_agent_for_plan_sketch = MagicMock() - policy = approach._plan_to_policy(plan, sketch=sketch) - state = _make_state() - policy(state) - monitor = _make_monitor(approach) - assert monitor.step(state) - # First divergence: budget 1 -> 0, replanned policy starts. - new_policy = approach._solve(Task(state, task.goal), timeout=10) - new_policy(state) - _sync(monitor, approach) - assert monitor.step(state) - # Second divergence: no budget left - the remaining plan resumes - # open-loop, with no further refinement attempt. - resumed = approach._solve(Task(state, task.goal), timeout=10) - approach._replan_suffix.assert_called_once() - approach._query_agent_for_plan_sketch.assert_not_called() - resumed(state) - status = approach.get_execution_monitoring_info()[0] - assert status.current_option.name == "Place" - - def test_reset_for_new_episode_clears_state(self): - """A new episode refreshes the budget and clears the live status.""" - approach, _, _ = _make_approach() - _enable_replanning(approach, 2) - assert approach._exec_replans_left == 2 - holding = {GroundAtom(_Holding, [_block0])} - plan, sketch = _make_two_step_plan(holding) - approach._plan_to_policy(plan, sketch=sketch) - assert approach.get_execution_monitoring_info() - approach._exec_replans_left = 0 - approach.reset_for_new_episode() - assert not approach.get_execution_monitoring_info() - assert approach._exec_replans_left == 2 - - def test_init_requires_subgoal_annotations_monitor(self): - """Enabling the budget without the monitor is a config error.""" - _, _, task = _make_approach() - utils.update_config({"agent_bilevel_max_execution_replans": 2}) - kwargs = dict( - initial_predicates=_ALL_PREDICATES, - initial_options=_ALL_OPTIONS, - types={_block_type, _robot_type}, - action_space=Box(low=-1, high=1, shape=(1, )), - train_tasks=[task], - option_model=MagicMock(), - ) - with pytest.raises(ValueError, match="subgoal_annotations"): - AgentModelBasedApproach(**kwargs) - utils.update_config({"execution_monitor": "subgoal_annotations"}) - AgentModelBasedApproach(**kwargs) - - def test_replan_suffix_walkback_and_validation(self): - """_replan_suffix tries the failed step first, walks back only to the - latest holding annotation, and forward-validates.""" - from predicators.agent_sdk import bilevel_sketch as bs - approach, _, task = _make_approach() - on_atom = {GroundAtom(_On, [_block0, _block1])} - holding = {GroundAtom(_Holding, [_block0])} - sketch = [ - _SketchStep(_PickDone, [_block0], on_atom), # holds (x close) - _SketchStep(_PickDone, [_block0], holding), # does not hold - _SketchStep(_PlaceDone, [_block0, _block1], holding), # failed - ] - # block0.x=0.5 == block1.x=0.5 so On holds; held=0 so Holding fails. - state = _make_state({_block0: [0.5, 0.2, 0.0]}) - tried = [] - - def _fake_refine(tsk, suffix, remaining, attempt=0): - del tsk, remaining, attempt # unused - tried.append(len(suffix)) - # Succeed only for the 2-step suffix (resume at step 1). - if len(suffix) == 2: - new_plan, _ = _make_two_step_plan(holding) - return new_plan, True - return [], False - - approach._refine_sketch = MagicMock(side_effect=_fake_refine) - with patch.object(bs, "validate_plan_forward", - return_value=(True, "")): - policy = approach._replan_suffix(state, task, sketch, 2, 10) - assert policy is not None - # Tried failed step (suffix len 1) first, then one step back - # (len 2); never walked past the holding annotation at step 0. - assert tried == [1, 2] - - -# --------------------------------------------------------------------------- -# Tests: scheduled-plans section in the solve/explore prompt -# --------------------------------------------------------------------------- - - -class TestScheduledPlansPromptSection: - """The explore prompt shows plans already generated this cycle so the next - request proposes a complementary plan instead of repeating the identical - one (run_20260707_112310 emitted the same 1-step plan for both of a cycle's - requests).""" - - @staticmethod - def _prompt(scheduled_plans): - from predicators.agent_sdk import sketch_prompts - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "seed": 42, - }) - state = _make_state() - task = Task(state, {GroundAtom(_On, [_block0, _block1])}) - return sketch_prompts.build_solve_prompt( - task, - all_predicates=_ALL_PREDICATES, - all_options=_ALL_OPTIONS, - scheduled_plans=scheduled_plans, - propose_params=True, - ) - - def test_section_absent_without_scheduled_plans(self): - """No scheduled-plans section is emitted when none were scheduled.""" - for empty in (None, []): - prompt = self._prompt(empty) - assert "Plans Already Scheduled This Cycle" not in prompt - - def test_section_lists_plans_and_asks_for_different_one(self): - """Scheduled plans are listed so the agent proposes a different one.""" - plans = [ - " 0: Pick(block0)[0.5000]", - " 0: Place(block0, block1)[0.1000, 0.2000]", - ] - prompt = self._prompt(plans) - assert "## Plans Already Scheduled This Cycle" in prompt - assert "Plan 1:\n 0: Pick(block0)[0.5000]" in prompt - assert "Plan 2:\n 0: Place(block0, block1)[0.1000, 0.2000]" in prompt - assert "data is complementary rather than redundant" in prompt - # The instruction must keep the request goal-directed (this is what - # preserves the train-solve early-stopping semantics). - assert "repeat the best plan" in " ".join(prompt.split()) - - -# --------------------------------------------------------------------------- -# Tests: turn-cap exhaustion handling in _solve -# --------------------------------------------------------------------------- - - -class TestTurnCapHandling: - """Hitting agent_sdk_max_agent_turns_per_iteration ends the attempt with a - best-effort submission instead of burning the sketch retries.""" - - @staticmethod - def _cap_result(subtype=None, num_turns=None): - return { - "type": "result", - "subtype": subtype, - "num_turns": num_turns, - "total_cost_usd": 1.0, - } - - def test_responses_hit_turn_cap(self): - """Cap detection: subtype is authoritative, num_turns is fallback.""" - approach, _, _ = _make_approach() - cap = approach._responses_hit_turn_cap - assert cap([self._cap_result(subtype="error_max_turns")]) - max_turns = 50 - utils.update_config( - {"agent_sdk_max_agent_turns_per_iteration": max_turns}) - assert cap([self._cap_result(subtype="success", num_turns=max_turns)]) - assert not cap( - [self._cap_result(subtype="success", num_turns=max_turns - 1)]) - assert not cap([{"type": "assistant", "content": []}]) - assert not cap([]) - - def test_sketch_query_records_turn_cap(self): - """A capped session with no final text still marks the cap before the - empty-plan-text failure propagates.""" - from predicators.approaches import ApproachFailure - approach, _, task = _make_approach() - responses = [self._cap_result(subtype="error_max_turns")] - with patch.object(approach, - '_query_agent_sync', - return_value=responses): - with pytest.raises(ApproachFailure, match="empty plan text"): - approach._query_agent_for_plan_sketch(task) - assert approach._last_sketch_query_hit_turn_cap - - def test_solve_one_query_per_attempt_on_turn_cap(self): - """A capped attempt takes the final nudge and stops.""" - from predicators.approaches import ApproachFailure - approach, _, task = _make_approach() - query = MagicMock( - return_value=[self._cap_result(subtype="error_max_turns")]) - nudge = MagicMock(return_value=None) - with patch.object(approach, '_query_agent_sync', query), \ - patch.object(approach, '_nudge_final_submission', nudge): - with pytest.raises(ApproachFailure, match="Bilevel solve failed"): - approach._solve(task, timeout=10) - assert query.call_count == 1 # no re-query on the same context - nudge.assert_called_once_with() - - def test_solve_restarts_with_no_nudge_until_final_attempt(self): - """Non-final attempts restart directly with no nudge. - - The best-effort submission nudge fires only on the FINAL - attempt, as the ultimate fallback. - """ - from predicators.approaches import ApproachFailure - approach, _, task = _make_approach() - utils.update_config({"agent_solve_max_attempts": 3}) - query = MagicMock( - return_value=[self._cap_result(subtype="error_max_turns")]) - nudge = MagicMock(return_value=None) - with patch.object(approach, '_query_agent_sync', query), \ - patch.object(approach, '_nudge_final_submission', nudge): - with pytest.raises(ApproachFailure, match="Bilevel solve failed"): - approach._solve(task, timeout=10) - assert query.call_count == 3 # one full query per attempt - nudge.assert_called_once_with() - - def test_solve_no_requery_on_non_cap_failure(self): - """A non-cap failure (e.g. unparseable output) ends the attempt too. - - The fresh-context restart is the ONLY retry: a query that merely - failed to submit does not buy a second full-price query on a - context that already contains whatever went wrong. - """ - from predicators.approaches import ApproachFailure - approach, _, task = _make_approach() - utils.update_config({"agent_solve_max_attempts": 3}) - # Well under the cap, but no plan text: a real error, not budget end. - query = MagicMock( - return_value=[self._cap_result(subtype="success", num_turns=5)]) - nudge = MagicMock(return_value=None) - with patch.object(approach, '_query_agent_sync', query), \ - patch.object(approach, '_nudge_final_submission', nudge): - with pytest.raises(ApproachFailure, match="Bilevel solve failed"): - approach._solve(task, timeout=10) - assert query.call_count == 3 # one per attempt, not one per query - nudge.assert_called_once_with() - - def test_attempt_end_reason_labels_journal_outcome(self): - """A capture-less attempt records WHY it ended. - - That reason is the one fact the next fresh-context attempt - cannot rediscover from the transcript it no longer has. - """ - from predicators.approaches import ApproachFailure - approach, _, task = _make_approach() - nudge = MagicMock(return_value=None) - - for subtype, num_turns, reason in [ - ("error_max_turns", None, "turn cap"), - ("success", 5, "no submission"), - ]: - query = MagicMock( - return_value=[self._cap_result(subtype, num_turns)]) - with patch.object(approach, '_query_agent_sync', query), \ - patch.object(approach, '_nudge_final_submission', nudge): - with pytest.raises(ApproachFailure, match=reason): - approach._solve(task, timeout=10) - assert approach._last_attempt_end_reason == reason - assert approach._attempt_outcome_text( - None, None) == f"no capture ({reason})" - - def test_nudge_best_effort_flag_set_and_cleared(self): - """The nudge exposes best-effort capture to the tools only for the - duration of its own query.""" - approach, _, _ = _make_approach() - seen = {} - - def _fake_query(message, **kwargs): - del kwargs # unused - seen["flag"] = approach._tool_context.capture_best_effort_plan - seen["message"] = message - return [] - - with patch.object(approach, '_query_agent_sync', _fake_query): - policy = approach._nudge_final_submission() - assert policy is None - assert seen["flag"] is True - assert "even if it does not fully reach the goal" in seen["message"] - assert not approach._tool_context.capture_best_effort_plan - - def test_nudge_returns_captured_best_effort_plan(self): - """The nudge consumes a captured plan into a policy even when the - rollout did not reach the goal.""" - approach, _, _ = _make_approach() - plan = [_Pick.ground([_block0], np.array([0.5], dtype=np.float32))] - sketch = [_SketchStep(_Pick, [_block0], None)] - - def _fake_query(message, **kwargs): - del message, kwargs # unused - # Simulate submit_plan's best-effort capture. - assert approach._tool_context.capture_best_effort_plan - approach._tool_context.solved_plan = plan - approach._tool_context.solved_sketch = sketch - approach._tool_context.solved_plan_reached_goal = False - return [] - - with patch.object(approach, '_query_agent_sync', _fake_query): - policy = approach._nudge_final_submission() - assert policy is not None - assert approach._tool_context.solved_plan is None - assert approach._tool_context.solved_plan_reached_goal is None - - -# --------------------------------------------------------------------------- -# Policy mode (agent_solve_policy_mode) -# --------------------------------------------------------------------------- - -_POLICY_SOURCE = ''' -def get_option(state, memory): - if memory.get("issued"): - return None - memory["issued"] = True - return "Pick(block0:block)[0.5]" -''' - - -def test_policy_mode_constructor_checks(): - """Policy mode rejects replans>0 and sim-free configs.""" - state = _make_state() - task = Task(state, {GroundAtom(_On, [_block0, _block1])}) - base_kwargs = dict(initial_predicates=_ALL_PREDICATES, - initial_options=_ALL_OPTIONS, - types={_block_type, _robot_type}, - action_space=Box(low=-1, high=1, shape=(1, )), - train_tasks=[task]) - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "seed": 42, - "agent_solve_policy_mode": True, - "agent_bilevel_max_execution_replans": 1, - "execution_monitor": "subgoal_annotations", - }) - with pytest.raises(ValueError, match="mutually exclusive"): - AgentModelBasedApproach(**base_kwargs) - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "seed": 42, - "agent_solve_policy_mode": True, - "agent_bilevel_max_execution_replans": 0, - "agent_planner_use_simulator": False, - }) - with pytest.raises(ValueError, match="use_simulator"): - AgentModelBasedApproach(**base_kwargs) - - -def test_consume_policy_capture_builds_executor(): - """A captured policy source composes and executes closed-loop.""" - approach, _, _ = _make_approach() - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "seed": 42, - "agent_solve_policy_mode": True, - "agent_bilevel_max_execution_replans": 0, - }) - state = _make_state() - task = Task(state, {GroundAtom(_On, [_block0, _block1])}) - approach._tool_context.current_task = task - approach._tool_context.solved_policy_source = _POLICY_SOURCE - approach._tool_context.solved_plan_reached_goal = True - policy = approach._consume_validated_plan() - assert policy is not None - info = approach._last_capture_info - assert info is not None and info.validated - assert "policy.py sha=" in info.plan_lines[0] - # The composed executor runs the issued option (terminal is always - # False here, so the first call returns that option's action). - action = policy(state) - assert isinstance(action, Action) - # No sketch monitor is armed in policy mode. - assert approach._exec_status is None - - -def test_execution_policy_surfaces_option_failures(): - """A failed option is surfaced to the policy, not episode-fatal.""" - from predicators.approaches import ApproachFailure - approach, _, _ = _make_approach() - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "seed": 42, - "agent_solve_policy_mode": True, - "agent_bilevel_max_execution_replans": 0, - "agent_policy_max_options": 10, - }) - not_initiable = ParameterizedOption( - "Broken", - types=[_block_type], - params_space=Box(low=np.array([0.0], dtype=np.float32), - high=np.array([1.0], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_false, - terminal=_always_false, - ) - seen = [] - - def option_fn(state, last_failure): - del state - seen.append(last_failure) - if last_failure is None and len(seen) == 1: - return not_initiable.ground([_block0], - np.array([0.5], dtype=np.float32)) - return None - - policy = approach._policy_to_execution_policy(option_fn) - with pytest.raises(ApproachFailure, match="DONE"): - policy(_make_state()) - assert seen[0] is None - assert seen[1] is not None # the failure was surfaced, not fatal - - -def test_execution_policy_budget_is_fatal(): - """The option cap converts an oscillating policy into a failure.""" - from predicators.approaches import ApproachFailure - approach, _, _ = _make_approach() - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "seed": 42, - "agent_solve_policy_mode": True, - "agent_bilevel_max_execution_replans": 0, - "agent_policy_max_options": 2, - }) - not_initiable = ParameterizedOption( - "Broken", - types=[_block_type], - params_space=Box(low=np.array([0.0], dtype=np.float32), - high=np.array([1.0], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_false, - terminal=_always_false, - ) - - def option_fn(state, last_failure): - del state, last_failure - return not_initiable.ground([_block0], np.array([0.5], - dtype=np.float32)) - - policy = approach._policy_to_execution_policy(option_fn) - with pytest.raises(ApproachFailure, match="option budget"): - policy(_make_state()) - - -def _broken_option(): - return ParameterizedOption( - "Broken", - types=[_block_type], - params_space=Box(low=np.array([0.0], dtype=np.float32), - high=np.array([1.0], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_false, - terminal=_always_false, - ) - - -def test_execution_policy_stuck_loop_is_fatal(): - """K consecutive failures of one identical command end the episode before - the option budget is burned (mirrors execute_policy_forward).""" - from predicators.approaches import ApproachFailure - approach, _, _ = _make_approach() - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "seed": 42, - "agent_solve_policy_mode": True, - "agent_bilevel_max_execution_replans": 0, - "agent_policy_max_options": 50, - }) - not_initiable = _broken_option() - issued = [] - - def option_fn(state, last_failure): - del state, last_failure - issued.append(1) - return not_initiable.ground([_block0], np.array([0.5], - dtype=np.float32)) - - policy = approach._policy_to_execution_policy(option_fn) - with pytest.raises(ApproachFailure, - match="re-issued the same failing option"): - policy(_make_state()) - assert len(issued) == 3 # the guard default, not the 50 cap - - -def test_execution_policy_stuck_loop_resets_on_changed_params(): - """Adapting the parameters after each failure avoids the guard: the episode - runs to the option budget instead.""" - from predicators.approaches import ApproachFailure - approach, _, _ = _make_approach() - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "seed": 42, - "agent_solve_policy_mode": True, - "agent_bilevel_max_execution_replans": 0, - "agent_policy_max_options": 6, - }) - not_initiable = _broken_option() - n_calls = [0] - - def option_fn(state, last_failure): - del state, last_failure - n_calls[0] += 1 - return not_initiable.ground([_block0], - np.array([0.1 * n_calls[0]], - dtype=np.float32)) - - policy = approach._policy_to_execution_policy(option_fn) - with pytest.raises(ApproachFailure, match="option budget"): - policy(_make_state()) - assert n_calls[0] == 6 - - -def test_execution_policy_stuck_loop_on_mid_execution_raise(): - """The guard also counts a skill that raises from inside its own policy - (e.g. a motion-planning refusal), whose exception carries no - last_failed_option of its own: option_policy_to_policy attributes it to the - executing option, so K identical re-issues end the episode instead of - burning the option budget (2026-08-25 policy-arm cycle-4 test).""" - from predicators.approaches import ApproachFailure - approach, _, _ = _make_approach() - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "seed": 42, - "agent_solve_policy_mode": True, - "agent_bilevel_max_execution_replans": 0, - "agent_policy_max_options": 50, - }) - - def _refusing_policy(_s, _m, _o, _p): - raise utils.OptionExecutionFailure( - "[Broken/MoveAbove] BiRRT collision: start configuration " - "in collision") - - refusing = ParameterizedOption( - "Broken", - types=[_block_type], - params_space=Box(low=np.array([0.0], dtype=np.float32), - high=np.array([1.0], dtype=np.float32)), - policy=_refusing_policy, - initiable=_always_true, - terminal=_always_false, - ) - issued = [] - - def option_fn(state, last_failure): - del state, last_failure - issued.append(1) - return refusing.ground([_block0], np.array([0.5], dtype=np.float32)) - - policy = approach._policy_to_execution_policy(option_fn) - with pytest.raises(ApproachFailure, - match="re-issued the same failing option"): - policy(_make_state()) - assert len(issued) == 3 # the guard default, not the 50 cap - - -def _instant_option(): - """Completes immediately (terminal always true) without failing.""" - return ParameterizedOption( - "Instant", - types=[_block_type], - params_space=Box(low=np.array([0.0], dtype=np.float32), - high=np.array([1.0], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_true, - terminal=_always_true, - ) - - -def test_execution_policy_noop_livelock_is_fatal(): - """K consecutive clean completions of one identical command with no - observable state change end the episode as a livelock (mirrors - execute_policy_forward).""" - from predicators.approaches import ApproachFailure - approach, _, _ = _make_approach() - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "seed": 42, - "agent_solve_policy_mode": True, - "agent_bilevel_max_execution_replans": 0, - "agent_policy_max_options": 50, - }) - instant = _instant_option() - issued = [] - - def option_fn(state, last_failure): - del state, last_failure - issued.append(1) - return instant.ground([_block0], np.array([0.5], dtype=np.float32)) - - policy = approach._policy_to_execution_policy(option_fn) - with pytest.raises(ApproachFailure, match="no observable state change"): - for _ in range(10): - policy(_make_state()) - assert len(issued) == 3 # the guard default, not the 50 cap - - -def test_execution_policy_noop_livelock_resets_on_changed_params(): - """Varying the parameters makes each command a different one, so the - livelock guard never trips; the option budget ends the episode.""" - from predicators.approaches import ApproachFailure - approach, _, _ = _make_approach() - utils.reset_config({ - "env": "cover", - "approach": "agent_model_based", - "seed": 42, - "agent_solve_policy_mode": True, - "agent_bilevel_max_execution_replans": 0, - "agent_policy_max_options": 6, - }) - instant = _instant_option() - n_calls = [0] - - def option_fn(state, last_failure): - del state, last_failure - n_calls[0] += 1 - return instant.ground([_block0], - np.array([0.1 * n_calls[0]], dtype=np.float32)) - - policy = approach._policy_to_execution_policy(option_fn) - with pytest.raises(ApproachFailure, match="option budget"): - for _ in range(10): - policy(_make_state()) - assert n_calls[0] == 6 diff --git a/tests/approaches/test_agent_nl_world_model_approach.py b/tests/approaches/test_agent_nl_world_model_approach.py deleted file mode 100644 index 3ed76b8215..0000000000 --- a/tests/approaches/test_agent_nl_world_model_approach.py +++ /dev/null @@ -1,181 +0,0 @@ -"""Tests for the natural-language world model approach's harness glue (paper -arm C3).""" -# pylint: disable=protected-access -import os -from types import SimpleNamespace -from typing import Any, List - -from predicators import utils -from predicators.agent_sdk import learn_prompts -from predicators.agent_sdk.tools import ToolContext -from predicators.approaches import agent_nl_world_model_approach as anl -from predicators.envs import create_new_env -from predicators.explorers.agent_model_free_explorer import \ - AgentModelFreeExplorer -from predicators.ground_truth_models import get_gt_options -from predicators.settings import CFG -from predicators.structs import Task - -_NOTES = "# Mechanisms\n- the hand moves to the PickPlace parameter\n" - - -def _cover() -> Any: - utils.reset_config({ - "env": "cover", - "num_train_tasks": 2, - "num_test_tasks": 1, - "agent_sdk_use_local_sandbox": True, - "seed": 0, - }) - env = create_new_env("cover") - train_tasks = [ - Task(t.task.init, t.task.goal, goal_nl="Cover the target.") - for t in env.get_train_tasks() - ] - options = get_gt_options(env.get_name()) - return env, train_tasks, options - - -def _bare(env: Any, train_tasks: List[Task], options: Any, - log_dir: str) -> Any: - approach = anl.AgentNotesWorldModelApproach.__new__( - anl.AgentNotesWorldModelApproach) - approach._types = env.types - approach._initial_predicates = set(env.predicates) - approach._initial_options = options - approach._train_tasks = train_tasks - approach._tool_context = ToolContext(types=env.types, - predicates=set(env.predicates), - options=options, - train_tasks=train_tasks) - approach._agent_session = None - approach._notes = "" - approach._notes_version = None - approach._online_learning_cycle = 0 - approach._get_log_dir = lambda: log_dir # type: ignore[method-assign] - approach._get_all_options = lambda: options # type: ignore[method-assign] - approach._get_all_trajectories = lambda: [] # type: ignore[method-assign] - approach._offline_dataset = SimpleNamespace( # type: ignore[assignment] - trajectories=[]) - approach._online_trajectories = [] - approach._option_model = None - approach._synthesized_samplers = {} - return approach - - -def test_predicate_allowlist_and_paths(tmp_path, monkeypatch) -> None: - """The kept-predicate allowlist applies; sandbox paths mirror the code - arms' mapping.""" - env, train_tasks, options = _cover() - approach = _bare(env, train_tasks, options, str(tmp_path)) - monkeypatch.setattr(CFG, "agent_sim_learn_kept_predicates_names", []) - assert approach._get_all_predicates() == set(env.predicates) - monkeypatch.setattr(CFG, "agent_sim_learn_kept_predicates_names", - ["Holding"]) - assert {p.name for p in approach._get_all_predicates()} == {"Holding"} - paths = approach._notes_paths() - assert paths["notes_file"] == os.path.join(str(tmp_path), "sandbox", - "world_model.md") - assert paths["notes_file_for_agent"] == "./world_model.md" - assert approach._get_synthesis_tool_names() == ["run_python"] - - -def test_learn_notes_gating_and_checkpoint_round_trip(tmp_path, - monkeypatch) -> None: - """No data means no session unless zero-shot is on; the document survives a - checkpoint and is written back into the sandbox.""" - env, train_tasks, options = _cover() - approach = _bare(env, train_tasks, options, str(tmp_path)) - calls: List[Any] = [] - approach._run_notes_session = calls.append - monkeypatch.setattr(CFG, "agent_sim_learn_zero_shot", False) - approach._learn_notes() - assert not calls - monkeypatch.setattr(CFG, "agent_sim_learn_zero_shot", True) - approach._learn_notes() - assert calls == [[]] - # A session's document is loaded from the sandbox file. - paths = approach._notes_paths() - os.makedirs(paths["base"], exist_ok=True) - with open(paths["notes_file"], "w", encoding="utf-8") as f: - f.write(_NOTES) - approach._load_notes(paths) - assert approach._notes == _NOTES - assert approach._notes_version is not None - assert approach._tool_context.world_model_notes == _NOTES - assert approach._tool_context.world_model_notes_path == \ - "./world_model.md" - # Checkpoint round trip into a fresh instance with an empty sandbox. - saved = approach._extra_save_state() - assert saved["world_model_notes"] == _NOTES - other_dir = os.path.join(str(tmp_path), "other") - other = _bare(env, train_tasks, options, other_dir) - other._load_extra_save_state(saved) - assert other._notes == _NOTES - with open(other._notes_paths()["notes_file"], encoding="utf-8") as f: - assert f.read() == _NOTES - - -def test_notes_reach_the_solve_and_explore_prompts(tmp_path) -> None: - """The document is quoted into the solve prompt, the explore prompt, and - the system prompt names it; without notes nothing is quoted.""" - env, train_tasks, options = _cover() - approach = _bare(env, train_tasks, options, str(tmp_path)) - approach._initial_image_section = lambda: "" # type: ignore[method-assign] - assert approach._solve_prompt_extra_sections() == "" - approach._notes = _NOTES - approach._sync_tool_context() - extra = approach._solve_prompt_extra_sections() - assert "World model notes (./world_model.md)" in extra - assert "hand moves to the PickPlace parameter" in extra - prompt = approach._build_solve_prompt(train_tasks[0]) - assert "hand moves to the PickPlace parameter" in prompt - assert prompt.index("World model notes") < prompt.index("## Objects") - system = approach._get_agent_system_prompt() - assert "world_model.md" in system and "no simulator" in system - approach._learning_mode = True - assert "natural-language document" in approach._get_agent_system_prompt() - approach._learning_mode = False - explorer = AgentModelFreeExplorer(set(env.predicates), options, env.types, - env.action_space, train_tasks, 10, - approach._tool_context, - None) # type: ignore[arg-type] - explore_prompt = explorer._build_exploration_prompt(0) - assert "hand moves to the PickPlace parameter" in explore_prompt - assert explore_prompt.index("World model notes") < \ - explore_prompt.index("## Instructions") - - -def test_notes_learn_message_and_exec_namespace(tmp_path) -> None: - """The first message carries the data roster, goals, and the prior-notes - pointer; the exec namespace exposes the data helpers.""" - env, train_tasks, options = _cover() - approach = _bare(env, train_tasks, options, str(tmp_path)) - paths = approach._notes_paths() - message = approach._build_notes_learn_message([], paths) - assert "0 recorded trajectories" in message - assert "Cover the target." in message - assert message.count("Cover the target.") == 1 - assert "Read it first" not in message - assert "./reference/structs.py" in message - assert os.path.isfile( - os.path.join(paths["base"], "reference", "structs.py")) - os.makedirs(paths["base"], exist_ok=True) - with open(paths["notes_file"], "w", encoding="utf-8") as f: - f.write(_NOTES) - assert "Read it first" in approach._build_notes_learn_message([], paths) - ns = approach._build_notes_exec_ns([]) - assert set(ns) >= { - "trajectories", "train_tasks", "is_goal_state", "describe_trajectory", - "np" - } - assert ns["is_goal_state"](train_tasks[0].init, 0) is False - - -def test_notes_prompt_builders() -> None: - """Builders render without leftovers; the block is empty without notes.""" - system = learn_prompts.build_notes_learn_system_prompt() - assert "# Mechanisms" in system and "__" not in system - assert learn_prompts.render_world_model_notes_block("", "x") == "" - zero = learn_prompts.render_notes_zero_shot_message() - assert "No trajectory has been recorded" in zero diff --git a/tests/approaches/test_agent_program_world_model_approach.py b/tests/approaches/test_agent_program_world_model_approach.py index 9c7234b8e2..c7f50261c0 100644 --- a/tests/approaches/test_agent_program_world_model_approach.py +++ b/tests/approaches/test_agent_program_world_model_approach.py @@ -15,7 +15,6 @@ from predicators.datasets import create_dataset from predicators.envs import create_new_env from predicators.ground_truth_models import get_gt_options -from predicators.settings import CFG _PROGRAM = ''' LATENT_FEATURES = {"robot": ["phase"]} @@ -69,9 +68,8 @@ def test_belief_particles_and_override_scope() -> None: latents all come from the installed program.""" env, train_tasks, options = _cover() approach = _bare(env, train_tasks, options) - # No model yet: no particles, and the initial latent is left alone. + # No model yet: no particles. assert not approach._belief_particles() - assert approach._attach_initial_latent(train_tasks[0]) is train_tasks[0] program, err = load_program_world_model(_PROGRAM, env.types, env.predicates, options) assert err is None and program is not None @@ -86,10 +84,6 @@ def test_belief_particles_and_override_scope() -> None: # The current task drives the draw when one is set. approach._tool_context.current_task = train_tasks[1] assert approach._belief_particles() == particles - # The nominal latent is attached to the task the planner sees. - task = approach._attach_initial_latent(train_tasks[0]) - assert task.init.latent is not None and "phase" in task.init.latent - assert train_tasks[0].init.latent is None # Under the scope every latent-less start rolls from the particle. model: ProgramOptionModel = approach._program_model (pick_place, ) = [o for o in options if o.name == "PickPlace"] @@ -108,32 +102,6 @@ def test_belief_particles_and_override_scope() -> None: assert approach._latent_tracking_available() is False -def test_learn_simulator_gating(monkeypatch) -> None: - """No data means no session unless zero-shot is on.""" - env, train_tasks, options = _cover() - approach = _bare(env, train_tasks, options) - approach._persist_fit_trajectories = lambda *a, **k: None - calls: List[Any] = [] - program, _ = load_program_world_model(_PROGRAM, env.types, env.predicates, - options) - - def _session(trajectories): - calls.append(list(trajectories)) - return program - - approach._run_program_synthesis_session = _session - monkeypatch.setattr(CFG, "agent_sim_learn_zero_shot", False) - approach._learn_simulator([]) - assert not calls and approach._program is None - monkeypatch.setattr(CFG, "agent_sim_learn_zero_shot", True) - approach._learn_simulator([]) - assert calls == [[]] and approach._program is program - dataset = create_dataset(env, train_tasks, options, env.predicates) - monkeypatch.setattr(CFG, "agent_sim_learn_zero_shot", False) - approach._learn_simulator(list(dataset.trajectories)) - assert len(calls) == 2 and len(calls[1]) == len(dataset.trajectories) - - def test_rehydrate_from_world_model_file(tmp_path, monkeypatch) -> None: """A checkpoint's world_model.py rebuilds the option model.""" env, train_tasks, options = _cover() diff --git a/tests/approaches/test_agent_sim_learning_ablations.py b/tests/approaches/test_agent_sim_learning_ablations.py index 8f9f9f1986..2c17d5ba08 100644 --- a/tests/approaches/test_agent_sim_learning_ablations.py +++ b/tests/approaches/test_agent_sim_learning_ablations.py @@ -5,8 +5,6 @@ ensemble. * ``agent_sim_learn_declared_params_only`` (A4): no estimation runs; the declaration is the estimate and its box the plausible interval. -* ``agent_sim_learn_zero_shot`` (A2): the synthesis session runs with - no recorded transitions. """ # pylint: disable=protected-access from typing import Any, Dict, List @@ -45,7 +43,6 @@ def _bare_approach() -> Any: approach._base_env = _RegistryEnv() approach._identified_physical_params = {} approach._identified_physical_sigma_points = [] - approach._cycle_applied_physical = {} approach._fitted_params = {} approach._param_ensemble = [] approach._param_specs = [] @@ -118,7 +115,6 @@ def test_deploy_declared_params_uses_the_declaration_as_the_estimate() -> None: # to the planning env. assert approach._fitted_params == {"k": 2.0} assert approach._base_env.applied == [{"lateral_friction": 0.5}] - assert approach._cycle_applied_physical == {"lateral_friction": 0.5} # Physics margin spans the declared box. frictions = [ p["lateral_friction"] @@ -209,7 +205,6 @@ def test_no_data_seeding_applies_declared_physical_inits() -> None: which the zero-shot arm relies on.""" utils.reset_config({ "agent_sim_learn_declared_params_only": False, - "agent_sim_learn_oracle_sim_params": False, "agent_explorer_info_seeking": False, }) approach = _bare_approach() @@ -220,40 +215,6 @@ def test_no_data_seeding_applies_declared_physical_inits() -> None: assert approach._last_fit_result is None -def test_zero_shot_flag_gates_data_free_synthesis() -> None: - """With no transitions, _learn_simulator returns early unless the zero-shot - flag is set, in which case synthesis runs on empty data.""" - approach: Any = asla.AgentSimLearningApproach.__new__( - asla.AgentSimLearningApproach) - approach._explainability_cache = {} - approach._sysid_fit_cache = {} - approach._persist_fit_trajectories = lambda *a, **k: None - approach._maybe_install_oracle_samplers = lambda: None - approach._extract_obs_triples = lambda trajs: [] - approach._residual_rules = None - approach._learned_simulator = None - approach._fitted_params = {} - calls: List[Any] = [] - - def _synth(trajectories, obs_triples, base_pred_triples, inferred_hint): - calls.append( - (trajectories, obs_triples, base_pred_triples, inferred_hint)) - - approach._synthesize_with_agent = _synth - utils.reset_config({ - "agent_sim_learn_zero_shot": False, - "agent_sim_learn_oracle_sim_program": False, - }) - approach._learn_simulator([]) - assert not calls - utils.reset_config({ - "agent_sim_learn_zero_shot": True, - "agent_sim_learn_oracle_sim_program": False, - }) - approach._learn_simulator([]) - assert calls == [([], [], [], {})] - - def test_declared_params_prompt_section_is_flag_gated() -> None: """The no-estimation section renders only under the A3 flag.""" kwargs: Dict[str, Any] = dict( diff --git a/tests/approaches/test_agent_sim_learning_approach.py b/tests/approaches/test_agent_sim_learning_approach.py index 397bdee1e6..5941407ec1 100644 --- a/tests/approaches/test_agent_sim_learning_approach.py +++ b/tests/approaches/test_agent_sim_learning_approach.py @@ -6,24 +6,20 @@ that solve a pybullet_boil task. """ # pylint: disable=protected-access -import inspect import logging import os import re from types import SimpleNamespace from typing import List, Optional, Sequence, Set, Tuple, cast -import dill as pkl import numpy as np import pytest from predicators import utils +from predicators.agent_sdk.sketch_types import SketchStep as _SketchStep from predicators.approaches import agent_sim_learning_approach as asla -from predicators.approaches.agent_model_based_approach import _SketchStep from predicators.approaches.agent_sim_learning_approach import \ AgentSimLearningApproach -from predicators.code_sim_learning.fit_space import FitResult -from predicators.code_sim_learning.identifiability import Verdict from predicators.code_sim_learning.utils import LearnedSimulator, \ apply_rules, merge_updates from predicators.envs import create_new_env @@ -595,60 +591,6 @@ def test_rollout_fit_trajectories_subset() -> None: obj._rollout_fit_trajectories(traj_idxs=[3]) -def _cross_cycle_fit(value: float) -> Tuple[FitResult, dict]: - result = FitResult(names=["friction"], - samples=np.array([[value]]), - log_probs=np.zeros(1), - jacobian=None, - noise_sigma=0.05, - prior_sigma=np.array([0.75]), - scales=["log"]) - report = { - "friction": { - "posterior_std": 0.1, - "prior_std": 0.75, - "contraction": 0.13, - "verdict": Verdict.IDENTIFIED, - "note": "", - } - } - return result, report - - -def test_cross_cycle_inconsistent_holds_then_confirms() -> None: - """A many-sigma jump is held once, accepted on independent repeat. - - Regression for run_20260724_232411 seed2: cycle fits 0.3236 -> - 0.6267 (4.7 combined sigmas). The first jump must flag INCONSISTENT - (trusted history unchanged, both fits recorded as hull candidates); - a following cycle re-fitting near the new value confirms the jump - and the history moves - without confirmation the stale reference - would flag every future fit forever. - """ - obj = object.__new__(AgentSimLearningApproach) - obj._sysid_fit_history = {} - obj._sysid_pending_fit = {} - - result, report = _cross_cycle_fit(0.3236) - obj._check_cross_cycle_consistency(result, report, ["friction"]) - assert report["friction"]["verdict"] is Verdict.IDENTIFIED - assert obj._sysid_fit_history["friction"][0] == 0.3236 - - result, report = _cross_cycle_fit(0.6267) - obj._check_cross_cycle_consistency(result, report, ["friction"]) - assert report["friction"]["verdict"] is Verdict.INCONSISTENT - assert report["friction"]["candidate_values"] == [0.3236, 0.6267] - # Trusted history holds; the rejected fit waits as pending. - assert obj._sysid_fit_history["friction"][0] == 0.3236 - assert obj._sysid_pending_fit["friction"][0] == 0.6267 - - result, report = _cross_cycle_fit(0.63) - obj._check_cross_cycle_consistency(result, report, ["friction"]) - assert report["friction"]["verdict"] is Verdict.IDENTIFIED - assert obj._sysid_fit_history["friction"][0] == 0.63 - assert "friction" not in obj._sysid_pending_fit - - def test_make_probe_process_model_factory() -> None: """The certificate-probe factory mirrors the combined simulator. @@ -705,124 +647,6 @@ def latent_rule(state: State, latent: dict, history: list, updates: dict, assert factory()(state, noop).get(thing, "x") == 1.0 -def test_cross_cycle_arbitration_by_pooled_evidence() -> None: - """A flagged jump is accepted when pooled data decisively backs it. - - Regression for run_20260727_210827 seed1: the sharp-but-biased - 2-trajectory cycle-0 fit (0.9313, true 0.5) was held over the - 4-trajectory refit (0.4748) for the rest of the run even though the - refit explained the pooled data ~30x better. With a pooled-SSE probe - the arbitration must accept the new value immediately; an ambivalent - gap (or a failing probe) must keep the hold. - """ - obj = object.__new__(AgentSimLearningApproach) - obj._sysid_fit_history = {} - obj._sysid_pending_fit = {} - - result, report = _cross_cycle_fit(0.9313) - obj._check_cross_cycle_consistency(result, report, ["friction"]) - assert obj._sysid_fit_history["friction"][0] == 0.9313 - - def pooled_sse(theta: dict) -> float: - return 0.14 if abs(theta["friction"] - 0.4748) < 1e-9 else 4.4 - - result, report = _cross_cycle_fit(0.4748) - obj._check_cross_cycle_consistency(result, - report, ["friction"], - pooled_sse=pooled_sse) - assert report["friction"]["verdict"] is Verdict.IDENTIFIED - assert "candidate_values" not in report["friction"] - assert obj._sysid_fit_history["friction"][0] == 0.4748 - assert "friction" not in obj._sysid_pending_fit - - # Ambivalent pooled gap (below the decisive ratio): hold as before. - obj._sysid_fit_history = {"friction": (0.9313, 0.1, "log")} - obj._sysid_pending_fit = {} - result, report = _cross_cycle_fit(0.4748) - obj._check_cross_cycle_consistency(result, - report, ["friction"], - pooled_sse=lambda theta: 0.14) - assert report["friction"]["verdict"] is Verdict.INCONSISTENT - assert obj._sysid_fit_history["friction"][0] == 0.9313 - assert obj._sysid_pending_fit["friction"][0] == 0.4748 - - # A failing SSE probe must fall back to the hold, not crash. - obj._sysid_fit_history = {"friction": (0.9313, 0.1, "log")} - obj._sysid_pending_fit = {} - - def broken_sse(theta: dict) -> float: - raise RuntimeError("env died") - - result, report = _cross_cycle_fit(0.4748) - obj._check_cross_cycle_consistency(result, - report, ["friction"], - pooled_sse=broken_sse) - assert report["friction"]["verdict"] is Verdict.INCONSISTENT - assert obj._sysid_fit_history["friction"][0] == 0.9313 - - -def test_persist_fit_trajectories(tmp_path, monkeypatch) -> None: - """Fit data lands in /fit_data/, one numbered pickle per fit.""" - obj = object.__new__(AgentSimLearningApproach) - obj._fit_trajectories = cast(List[LowLevelTrajectory], - ["fake_traj_a", "fake_traj_b"]) - obj._physical_param_specs = [] - obj._identified_physical_params = {"friction": 0.5} - obj._get_log_dir = lambda: str(tmp_path) # type: ignore[method-assign] - - obj._persist_fit_trajectories() - obj._persist_fit_trajectories() - out_dir = tmp_path / "fit_data" - files = sorted(f.name for f in out_dir.glob("*.pkl")) - assert files == [ - "fit_trajectories_000_fitted.pkl", "fit_trajectories_001_fitted.pkl" - ] - with open(out_dir / files[0], "rb") as f: - payload = pkl.load(f) - assert payload["trajectories"] == ["fake_traj_a", "fake_traj_b"] - assert payload["identified_physical_params"] == {"friction": 0.5} - - monkeypatch.setattr(CFG, "code_sim_learning_persist_fit_data", False) - obj._persist_fit_trajectories() - assert len(list(out_dir.glob("*.pkl"))) == 2 - - -def test_fit_data_is_dumped_even_when_no_fit_runs(tmp_path) -> None: - """A cycle that declines to fit is exactly the one worth post-morteming. - - Persistence used to sit only inside the sysID fit, so the branch - that never ran was the branch whose data mattered: - run_20260817_171402 declined on a sweep returning one identical SSE - for every value of five parameters, and left nothing on disk to - explain it. Dumping where the data ARRIVES is what makes that - replayable. - """ - obj = object.__new__(AgentSimLearningApproach) - obj._physical_param_specs = [] - obj._identified_physical_params = {} - obj._get_log_dir = lambda: str(tmp_path) # type: ignore[method-assign] - obj._explainability_cache = {} - obj._sysid_fit_cache = {} - - # The one line of _learn_simulator this is about, with no fit after it. - obj._fit_trajectories = cast(List[LowLevelTrajectory], ["traj"]) - obj._persist_fit_trajectories("recorded") - - files = [f.name for f in (tmp_path / "fit_data").glob("*.pkl")] - assert files == ["fit_trajectories_000_recorded.pkl"] - with open(tmp_path / "fit_data" / files[0], "rb") as f: - assert pkl.load(f)["trajectories"] == ["traj"] - - # The wiring, not just the function: _learn_simulator runs on every - # cycle whether or not a fit follows, so the dump has to hang off it. - # Asserted on the source because calling _learn_simulator for real - # needs a whole synthesis session, and without this the test above - # passes with the call deleted. - source = inspect.getsource(AgentSimLearningApproach._learn_simulator) # pylint: disable=protected-access - assert "_persist_fit_trajectories(\"recorded\")" in source, \ - "the recorded-data dump is no longer wired into _learn_simulator" - - def test_base_sim_reference_provisioning() -> None: """Base-sim source rides the sandbox reference registry (so every session phase gets it) and the agent-visible paths map per backend.""" @@ -900,11 +724,9 @@ def _make_checkpoint_stub(tmp_path, monkeypatch): obj._carried_physical_prior = {"lateral_friction": 0.5} obj._fit_evidence_history = {"vers_001": {"log_evidence": -1.0}} obj._identified_physical_sigma_points = [{"lateral_friction": 0.55}] - obj._sysid_fit_history = {} obj._residual_features = {"block": ["x"]} obj._current_simulator_version = "cycle_001_vers_003" obj._current_predicates_version = None - obj._current_samplers_version = None return obj, sandbox @@ -987,7 +809,6 @@ def test_rehydrate_rebuilds_simulator_from_restored_file( obj._learned_simulator = None obj._latent_init = None obj._fit_trajectories = [] - obj._synthesized_samplers = {} obj._base_env = SimpleNamespace(get_physical_param_info=lambda: {}) calls = [] monkeypatch.setattr( @@ -1008,8 +829,6 @@ def test_rehydrate_rebuilds_simulator_from_restored_file( monkeypatch.setattr(AgentSimLearningApproach, "_apply_identified_physical_params", lambda self, p: calls.append(("apply", dict(p)))) - monkeypatch.setattr(AgentSimLearningApproach, "_samplers_enabled", - staticmethod(lambda: False)) monkeypatch.setattr(AgentSimLearningApproach, "_rebuild_param_ensemble", lambda self: calls.append("ensemble")) # Checkpointed fitted params carry a stale name -> fall back to init. @@ -1040,39 +859,3 @@ def test_rehydrate_without_simulator_is_graceful(tmp_path, monkeypatch): lambda self, base: hooks.append(base)) obj._rehydrate_from_artifacts() assert hooks == [str(sandbox)] - - -def test_checkpoint_cycle_counter_semantics(tmp_path, monkeypatch): - """``_c`` files store c even when saved after the counter advanced, and - loading ``_None`` resumes at cycle 0 while ``_c`` resumes at c+1.""" - # pylint: disable=import-outside-toplevel - from predicators.approaches.agent_model_free_approach import \ - AgentModelFreeApproach - from predicators.structs import Dataset - utils.reset_config({ - "env": "cover", - "approach": "agent_model_free", - "seed": 0, - "approach_dir": str(tmp_path), - }) - obj = object.__new__(AgentModelFreeApproach) - obj._offline_dataset = Dataset([]) - obj._online_trajectories = [] - obj._run_id = "run" - obj._agent_session = None - monkeypatch.setattr(AgentModelFreeApproach, "_sync_tool_context", - lambda self: None) - # Post-offline checkpoint: counter 0, file _None. - obj._online_learning_cycle = 0 - obj.save(None) - # Cycle-3 checkpoint written AFTER the counter already advanced to 4 - # (the sim-learning subclass saves post-learn): the file must still - # denote cycle 3. - obj._online_learning_cycle = 4 - obj.save(3) - fresh = object.__new__(AgentModelFreeApproach) - fresh._agent_session = None - fresh.load(None) - assert fresh._online_learning_cycle == 0 - fresh.load(3) - assert fresh._online_learning_cycle == 4 diff --git a/tests/approaches/test_agent_sim_prompt_formatting.py b/tests/approaches/test_agent_sim_prompt_formatting.py index 80971abb8b..9b81853698 100644 --- a/tests/approaches/test_agent_sim_prompt_formatting.py +++ b/tests/approaches/test_agent_sim_prompt_formatting.py @@ -1,11 +1,5 @@ -"""Tests for synthesis-prompt formatter helpers. - -These are pure-Python staticmethods (or `self`-less methods) on -``AgentSimLearningApproach`` and ``AgentSimPredicateInventionApproach`` -that render parts of the agent's first synthesis message. They were -added so the agent (a) knows the provenance of each interaction -trajectory and (b) gets reminded about prior-cycle files in the sandbox. -""" +"""Tests for the ``evaluate_trajectory`` helper the synthesis namespace offers, +and for the learn system prompt's deliverables.""" # pylint: disable=protected-access,import-outside-toplevel,unused-import from __future__ import annotations @@ -25,273 +19,18 @@ def _approach_cls(): return AgentSimLearningApproach -def _mk_traj(is_demo, - task_idx, - sim_v=None, - preds_v=None, - reward=None, - terminated=None): - """Build a 1-action trajectory with the given provenance tags.""" - cup_type = Type("cup_type", ["f"]) - cup = cup_type("cup") - states = [State({cup: [0.0]}), State({cup: [1.0]})] - actions = [Action(np.array([0.5]))] - return LowLevelTrajectory( - states, - actions, - _is_demo=is_demo, - _train_task_idx=task_idx, - _source_simulator_version=sim_v, - _source_predicates_version=preds_v, - _env_reward=reward, - _env_terminated=terminated, - ) - - # ── _format_trajectory_listing ────────────────────────────────────── - -def test_trajectory_listing_empty(approach_cls): - """Empty trajectory list short-circuits to an empty string.""" - assert approach_cls._format_trajectory_listing([]) == "" - - -def test_trajectory_listing_demo_has_no_provenance_tail(approach_cls): - """Demo trajectories never carry provenance - even if the tags are set, the - listing should still render them as plain demos for consistency with the - offline-data semantics.""" - trajs = [_mk_traj(is_demo=True, task_idx=0)] - out = approach_cls._format_trajectory_listing(trajs) - assert "[0] demo, task 0" in out - assert "generated using" not in out - - -def test_trajectory_listing_interaction_with_provenance(approach_cls): - """Interaction trajectories with provenance show the sim/preds tags.""" - trajs = [ - _mk_traj(is_demo=False, - task_idx=2, - sim_v="cycle_001_vers_004", - preds_v="cycle_001_vers_003"), - ] - out = approach_cls._format_trajectory_listing(trajs) - assert "[0] interaction, task 2" in out - assert "sim cycle_001_vers_004" in out - assert "predicates cycle_001_vers_003" in out - - -def test_trajectory_listing_supervisor_rejected(approach_cls): - """A rejected episode surfaces only through its (reward, terminated) - pair: terminated with solved=0 and no bonus in the reward. No - REJECTED flag or violation specifics reach the roster - the rules - live in the NL goal description, so the agent must infer the - violation from its own trajectory rather than be told it. - """ - trajs = [ - _mk_traj(is_demo=False, task_idx=0), - _mk_traj(is_demo=False, task_idx=3, reward=-0.05, terminated=True), - ] - out = approach_cls._format_trajectory_listing(trajs) - lines = [l for l in out.splitlines() if l.startswith(" [")] - assert "REJECTED" not in out - assert "env reward=-0.05 (solved=0)" in lines[1] - # No violation specifics leak into the roster line. - assert "domino" not in lines[1] - assert "push" not in lines[1].lower() - - -def test_trajectory_listing_env_reward(approach_cls): - """Evaluated episodes show the env reward with a success flag; a rejected - topple counts as solved=0 even though it terminated.""" - trajs = [ - _mk_traj(is_demo=False, task_idx=0, reward=0.85, terminated=True), - _mk_traj(is_demo=False, task_idx=1, reward=-0.05, terminated=True), - _mk_traj(is_demo=False, task_idx=2), # never evaluated - ] - out = approach_cls._format_trajectory_listing(trajs) - lines = [l for l in out.splitlines() if l.startswith(" [")] - assert "env reward=0.85 (solved=1)" in lines[0] - assert "env reward=-0.05 (solved=0)" in lines[1] - assert "REJECTED" not in lines[1] - assert "reward" not in lines[2] - - -def test_trajectory_listing_partial_provenance(approach_cls): - """Only ``source_simulator_version`` set: list just the sim tag. - - No stray ``, `` may appear from the missing predicate half of the - provenance pair. - """ - trajs = [_mk_traj(is_demo=False, task_idx=1, sim_v="cycle_001_vers_007")] - out = approach_cls._format_trajectory_listing(trajs) - line = [l for l in out.splitlines() if l.startswith(" [0]")][0] - assert "sim cycle_001_vers_007" in line - assert "predicates" not in line - - # ── _format_prior_state_block ──────────────────────────────────────── - -def test_prior_state_block_empty_when_no_files(approach_cls, tmp_path): - """Neither simulator.py nor predicates.py exists → empty block.""" - out = approach_cls._format_prior_state_block(None, str(tmp_path)) - assert out == "" - - -def test_prior_state_block_simulator_only(approach_cls, tmp_path): - """Only simulator.py exists → block mentions it and not predicates.py.""" - (tmp_path / "simulator.py").write_text("# sim") - out = approach_cls._format_prior_state_block(None, str(tmp_path)) - assert "`./simulator.py`" in out - assert "`./predicates.py`" not in out - # Always points at the versioned-snapshot dirs for cross-reference. - assert "./simulator_versions/" in out - - -def test_prior_state_block_both_files(approach_cls, tmp_path): - """Both files exist → block lists them joined with ' and '.""" - (tmp_path / "simulator.py").write_text("# sim") - (tmp_path / "predicates.py").write_text("LEARNED_PREDICATES = []") - out = approach_cls._format_prior_state_block(None, str(tmp_path)) - assert "`./simulator.py` and `./predicates.py`" in out - # Soft language so the agent isn't forbidden from a fresh rewrite. - assert "fresh rewrite is fine" in out - - # ── _format_goal_nl_block (predicate-invention subclass) ──────────── - -def test_goal_nl_block_empty_when_no_tasks_have_goal_nl(): - """No ``goal_nl`` populated → empty block (no header).""" - from predicators.approaches.agent_sim_predicate_invention_approach import \ - AgentSimPredicateInventionApproach - fake_self = type( - "_FakeApproach", - (), - { - "_train_tasks": [type("_T", (), {"goal_nl": None})()] * 2, - }, - )() - out = AgentSimPredicateInventionApproach._format_goal_nl_block(fake_self) - assert out == "" - - -def test_goal_nl_block_dedups_identical_goals(): - """Same NL goal across tasks shows up once, with the single-task header.""" - from predicators.approaches.agent_sim_predicate_invention_approach import \ - AgentSimPredicateInventionApproach - fake_task = type("_T", (), {"goal_nl": "boil the water"}) - fake_self = type( - "_FakeApproach", - (), - { - "_train_tasks": [fake_task() for _ in range(3)], - }, - )() - out = AgentSimPredicateInventionApproach._format_goal_nl_block(fake_self) - assert out.startswith("Goal (natural language): boil the water") - # Trailing blank line separates from the next paragraph in the prompt. - assert out.endswith("\n\n") - - -def test_goal_nl_block_multiple_distinct_goals(): - """Distinct goals across tasks render as a bulleted list.""" - from predicators.approaches.agent_sim_predicate_invention_approach import \ - AgentSimPredicateInventionApproach - tasks = [ - type("_T1", (), {"goal_nl": "boil the water"})(), - type("_T2", (), {"goal_nl": "stack the cups"})(), - ] - fake_self = type("_FakeApproach", (), {"_train_tasks": tasks})() - out = AgentSimPredicateInventionApproach._format_goal_nl_block(fake_self) - assert "Goals across train tasks (natural language):" in out - assert " - boil the water" in out - assert " - stack the cups" in out - - # ── _build_synthesis_system_prompt (FO vs PO rule signature) ──────── # These render the whole synthesis system prompt. The method only touches # ``self`` through pure no-state helpers (``_rule_signature_section``, # ``_residual_rule_signature``, ``_extra_synthesis_system_prompt``), so a # bare instance via ``object.__new__`` is enough to render it. - -def _render_prompt(cls): - from predicators.approaches.agent_sim_learning_approach import \ - AgentSimLearningApproach - return AgentSimLearningApproach._build_synthesis_system_prompt( - object.__new__(cls)) - - -def test_synthesis_prompt_no_leftover_placeholders(approach_cls): - """Every templated placeholder is substituted in the rendered prompt.""" - prompt = _render_prompt(approach_cls) - for placeholder in ("__RULE_SIGNATURE_SECTION__", - "__RESIDUAL_RULE_SIGNATURE__", - "__SYNTHESIS_PROMPT_EXTRA__"): - assert placeholder not in prompt - - -def test_synthesis_prompt_sections_not_duplicated(approach_cls): - """The system prompt has exactly one of each major section. - - Guards against the bad-merge artifact that duplicated the Tools / - Refinement / Plan-format blocks (and double-injected the extra). - """ - prompt = _render_prompt(approach_cls) - for header in ("## `simulator.py`: a simulator subclass", - "## Step and restoration behavior", "## Plan format", - "## Fit and validate complete rollouts"): - assert prompt.count(header) == 1, (header, prompt.count(header)) - - -def test_fo_prompt_uses_subclass_contract(approach_cls): - """Fully observable models use the same subclass contract as PO ones.""" - prompt = _render_prompt(approach_cls) - assert "class MyDynamics(BaseSimulator):" in prompt - assert "RESIDUAL_ENV = MyDynamics" in prompt - assert "def _domain_specific_step(self):" in prompt - assert "RESIDUAL_RULES" not in prompt - assert "def residual_rule(" not in prompt - assert "## Hidden model state" not in prompt - - -def test_po_prompt_uses_subclass_memory_contract(): - """Both PO approaches receive one canonical model-state callback. - - A competing rule signature must not reappear beside the subclass - contract, and only predicate invention adds classifier guidance. - """ - import re - - from predicators.approaches.agent_sim_learning_approach import \ - AgentSimLearningApproach - from predicators.approaches.agent_sim_predicate_invention_approach import \ - AgentSimPredicateInventionApproach - from predicators.settings import CFG - old_flag = CFG.partially_observable - CFG.partially_observable = True - try: - for cls in (AgentSimLearningApproach, - AgentSimPredicateInventionApproach): - prompt = _render_prompt(cls) - assert "class MyDynamics(BaseSimulator):" in prompt - assert ("def update_model_state(cls, observation, model_state, " - "params, action):" in prompt) - assert "RESIDUAL_RULES" not in prompt - assert "LATENT_INIT" not in prompt - assert "def residual_rule(" not in prompt - # Memory guidance is injected exactly once. - headers = re.findall(r"(?m)^## Hidden model state$", prompt) - assert len(headers) == 1, cls - # The predicate-side latent guidance is invention-only. - has_pred_section = "### Predicate signature" in prompt - assert has_pred_section == ( - cls is AgentSimPredicateInventionApproach), cls - finally: - CFG.partially_observable = old_flag - - # ── _make_evaluate_trajectory_fn / _format_objective_block ────────── @@ -420,38 +159,6 @@ def _scope(physical_overrides=None): assert fn(states, None, task_idx=0, physics_sweep=True)["sweep"] is None -def test_format_objective_block(approach_cls): - """The objective block renders the first stated objective and is empty when - no evaluator states one.""" - from types import SimpleNamespace - - from predicators.structs import TaskEvaluator - - class _StatingEvaluator(TaskEvaluator): - """Evaluator with a public objective statement.""" - - def objective_description(self): - return "Topple the target legitimately; each blue costs 0.05." - - cup_type = Type("cup_type", ["f"]) - init = State({cup_type("cup"): [0.0]}) - - def _task(evaluator=None): - return Task(init, set(), evaluator=evaluator) - - fmt = approach_cls._format_objective_block - assert fmt(SimpleNamespace(_train_tasks=[])) == "" - assert fmt(SimpleNamespace(_train_tasks=[_task()])) == "" - assert fmt( - SimpleNamespace(_train_tasks=[_task(TaskEvaluator(set()))])) == "" - out = fmt( - SimpleNamespace( - _train_tasks=[_task(), _task(_StatingEvaluator(set()))])) - assert "## Task objective (env ground-truth reward)" in out - assert "each blue costs 0.05" in out - assert "evaluate_trajectory" in out - - def test_learn_message_ships_goal_required_mechanisms_as_hypotheses(): """The learn message distinguishes a hypothesis the goal can do without (record, do not ship) from one the goal REQUIRES (ship as a labelled diff --git a/tests/approaches/test_agent_solve_restart.py b/tests/approaches/test_agent_solve_restart.py deleted file mode 100644 index 682853c753..0000000000 --- a/tests/approaches/test_agent_solve_restart.py +++ /dev/null @@ -1,459 +0,0 @@ -"""Tests for the time-boxed restart loop in AgentModelBasedApproach._solve. - -The loop runs up to ``agent_solve_max_attempts`` solve attempts, each on -a fresh conversation when ``agent_solve_fresh_context`` is set: a -validated (evaluator-solved) capture returns immediately, best-effort -captures are banked and ranked by evaluator reward, journal auto-entries -record every attempt, and total failure re-raises the last error. -""" -# pylint: disable=protected-access -import time - -import numpy as np -import pytest -from gym.spaces import Box - -from predicators import utils -from predicators.agent_sdk import journal as journal_mod -from predicators.agent_sdk.session_base import AgentSessionFatalError -from predicators.approaches import ApproachFailure, ApproachTimeout -from predicators.approaches.agent_model_based_approach import \ - AgentModelBasedApproach, _CaptureInfo -from predicators.structs import Action, GroundAtom, Object, \ - ParameterizedOption, Predicate, State, Task, Type - -_block_type = Type("block", ["x"]) -_block0 = Object("block0", _block_type) -_Reached = Predicate("Reached", [_block_type], - lambda s, o: s.get(o[0], "x") >= 0.9) - - -def _noop_policy(_s, _m, _o, _p): - return Action(np.zeros(1, dtype=np.float32)) - - -_Move = ParameterizedOption( - "Move", - types=[_block_type], - params_space=Box(low=np.array([0.0], dtype=np.float32), - high=np.array([1.0], dtype=np.float32)), - policy=_noop_policy, - initiable=lambda _s, _m, _o, _p: True, - terminal=lambda _s, _m, _o, _p: False, -) - - -def _make_approach(overrides=None, sandbox_dir=None): - state = State({_block0: np.array([0.0], dtype=np.float32)}) - task = Task(state, {GroundAtom(_Reached, [_block0])}) - config = { - "env": "cover", - "approach": "agent_model_based", - "num_train_tasks": 1, - "num_test_tasks": 1, - "seed": 42, - } - config.update(overrides or {}) - utils.reset_config(config) - approach = AgentModelBasedApproach( - initial_predicates={_Reached}, - initial_options={_Move}, - types={_block_type}, - action_space=Box(low=-1, high=1, shape=(1, )), - train_tasks=[task], - option_model=None, - ) - if sandbox_dir is not None: - approach._tool_context.sandbox_dir = sandbox_dir - return approach, task - - -class _AttemptScript: - """Scripted per-attempt outcomes for a stubbed _solve_attempt. - - Each script item is ``("validated", reward)``, ``("best_effort", - reward)``, ``("fail", message)`` (raises ApproachFailure), or - ``("error", exception)`` (raises that exception). The stub sets - ``_last_capture_info`` exactly as _consume_validated_plan would. - """ - - def __init__(self, approach, script): - self.approach = approach - self.script = list(script) - self.calls = 0 - self.policies = [] - - def __call__(self, _task): - kind, value = self.script[self.calls] - self.calls += 1 - if kind == "fail": - raise ApproachFailure(value) - if kind == "error": - raise value - self.approach._last_capture_info = _CaptureInfo( - validated=(kind == "validated"), - reward=value, - plan_lines=[f"Move(block0:block)[{0.9 + self.calls / 100.0}]"], - validation_summary="validation: 5/6 rollouts ok") - policy = lambda _s: Action(np.zeros(1, dtype=np.float32)) - self.policies.append(policy) - return policy - - -def test_validated_capture_returns_immediately(): - """A validated first attempt short-circuits the loop.""" - approach, task = _make_approach({ - "agent_solve_max_attempts": 3, - "agent_solve_fresh_context": True, - }) - closes = [] - approach._close_agent_session = lambda: closes.append(1) - script = _AttemptScript(approach, [("validated", 0.95)]) - approach._solve_attempt = script - policy = approach._solve(task, timeout=10) - assert policy is script.policies[0] - assert script.calls == 1 - assert len(closes) == 1 # fresh context for the (only) attempt - - -def test_best_effort_banked_and_best_reward_wins(): - """Best-effort captures are ranked by evaluator reward across attempts.""" - approach, task = _make_approach({"agent_solve_max_attempts": 3}) - script = _AttemptScript(approach, [ - ("best_effort", -0.10), - ("best_effort", 0.40), - ("fail", "attempt 3 found nothing"), - ]) - approach._solve_attempt = script - policy = approach._solve(task, timeout=10) - assert script.calls == 3 - assert policy is script.policies[1] # the reward-0.40 capture - - -def test_validated_on_later_attempt_beats_banked_best_effort(): - """A later validated capture wins over an earlier banked one.""" - approach, task = _make_approach({"agent_solve_max_attempts": 3}) - script = _AttemptScript(approach, [ - ("best_effort", 0.90), - ("validated", 0.95), - ]) - approach._solve_attempt = script - policy = approach._solve(task, timeout=10) - assert script.calls == 2 - assert policy is script.policies[1] - - -def test_all_attempts_fail_reraises(): - """With nothing captured anywhere, the last failure propagates.""" - approach, task = _make_approach({"agent_solve_max_attempts": 2}) - script = _AttemptScript(approach, [ - ("fail", "first"), - ("fail", "second"), - ]) - approach._solve_attempt = script - with pytest.raises(ApproachFailure, match="second"): - approach._solve(task, timeout=10) - assert script.calls == 2 - - -def test_no_fresh_context_keeps_session(): - """agent_solve_fresh_context=False never closes the session.""" - approach, task = _make_approach({ - "agent_solve_max_attempts": 2, - "agent_solve_fresh_context": False, - }) - closes = [] - approach._close_agent_session = lambda: closes.append(1) - script = _AttemptScript(approach, [("fail", "a"), ("validated", 0.9)]) - approach._solve_attempt = script - approach._solve(task, timeout=10) - assert not closes - - -def test_journal_auto_entries_record_each_attempt(tmp_path): - """Every attempt leaves an auto entry with outcome and plan.""" - approach, task = _make_approach( - { - "agent_solve_max_attempts": 2, - "agent_solve_use_journal": True, - }, - sandbox_dir=str(tmp_path)) - approach._tool_context.test_task_idx = 0 - script = _AttemptScript(approach, [ - ("fail", "nothing"), - ("best_effort", -0.05), - ]) - approach._solve_attempt = script - approach._solve(task, timeout=10) - content = journal_mod.read_journal(str(tmp_path), - filename=journal_mod.ATTEMPTS_FILENAME) - assert "### cycle 0 task 0 attempt 1/2 (auto)" in content - assert "- outcome: no capture" in content - assert "### cycle 0 task 0 attempt 2/2 (auto)" in content - assert "best-effort capture (evaluator reward -0.05)" in content - assert "Move(block0:block)" in content - # The capture-time validation record travels into the auto entry so a - # later fresh-context attempt sees how reliable the capture was. - assert "- validation: 5/6 rollouts ok" in content - # Task context (goal + init state dict, the prompt's own - # representation) is a dedicated entry written once, at the TOP of - # the task's section - before any attempt entry. - assert content.count("### cycle 0 task 0 goal + initial state (auto)") == 1 - assert content.count("- goal: Reached(block0:block)") == 1 - assert content.count("- initial state features:") == 1 - assert "'block0:block'" in content - assert "'x'" in content - assert content.index("goal + initial state") < content.index( - "### cycle 0 task 0 attempt 1/2") - - -def test_attempt_bookkeeping_reset_per_attempt(): - """attempt_index/rollout counter/deadline are set and cleared.""" - approach, task = _make_approach({ - "agent_solve_max_attempts": 2, - "agent_solve_attempt_wall_clock": 2700, - }) - ctx = approach._tool_context - seen = [] - - def _attempt(_task): - seen.append( - (ctx.attempt_index, ctx.attempt_rollout_count, ctx.attempt_deadline - is not None)) - ctx.attempt_rollout_count += 7 - raise ApproachFailure("no capture") - - approach._solve_attempt = _attempt - with pytest.raises(ApproachFailure): - approach._solve(task, timeout=10) - assert seen == [(1, 0, True), (2, 0, True)] - assert ctx.attempt_deadline is None - assert ctx.attempt_start is None - # attempt_index resets too: a stale index would mislabel journal - # entries recorded outside any attempt. - assert ctx.attempt_index == 0 - - -def test_attempt_wall_spent(): - """_attempt_wall_spent trips BEFORE the deadline, at the spent floor. - - Agents that watch the [budget] footer end their query shortly before - the deadline (run_20260718_125643 queries 002-003); an attempt whose - remaining tools would only refuse must be labelled spent, not "no - submission". - """ - approach, _task = _make_approach({"agent_solve_attempt_wall_clock": 2700}) - ctx = approach._tool_context - # No deadline set: never spent. - assert not approach._attempt_wall_spent() - ctx.attempt_deadline = time.monotonic() - 1.0 - assert approach._attempt_wall_spent() - # 2 min left of 45: below the 20% floor (540s), counts as spent. - ctx.attempt_deadline = time.monotonic() + 120.0 - assert approach._attempt_wall_spent() - # 20 min left: plenty of budget still on the clock. - ctx.attempt_deadline = time.monotonic() + 1200.0 - assert not approach._attempt_wall_spent() - - -def test_nudge_suspends_and_restores_attempt_deadline(): - """The nudge must not permanently disarm the wall clock. - - It suspends the deadline so its own submission is not refused, but - a helper that leaked the None would let any future mid-attempt - caller run unbounded - the exact runaway the time-box targets. - """ - approach, _task = _make_approach({}) - ctx = approach._tool_context - deadline = time.monotonic() + 1000.0 - ctx.attempt_deadline = deadline - approach._query_agent_sync = lambda *a, **k: [] - policy = approach._nudge_final_submission() - assert policy is None - assert ctx.attempt_deadline == deadline - assert ctx.capture_best_effort_plan is False - - -def test_unexpected_error_salvages_banked_capture(): - """A non-ApproachFailure error executes the banked best-effort capture - instead of forfeiting it, and stops further attempts.""" - approach, task = _make_approach({"agent_solve_max_attempts": 3}) - ctx = approach._tool_context - script = _AttemptScript(approach, [ - ("best_effort", 0.40), - ("error", RuntimeError("pybullet exploded")), - ]) - approach._solve_attempt = script - policy = approach._solve(task, timeout=10) - assert policy is script.policies[0] - assert script.calls == 2 # attempt 3 never runs - # Bookkeeping fully cleaned despite the unexpected error. - assert ctx.attempt_start is None - assert ctx.attempt_deadline is None - assert ctx.attempt_index == 0 - - -def test_fatal_session_error_reraises_even_with_banked_capture(): - """AgentSessionFatalError is never salvaged by a banked capture: the - session backend is unusable, so the run must terminate (bookkeeping still - cleaned by the finally).""" - approach, task = _make_approach({"agent_solve_max_attempts": 3}) - ctx = approach._tool_context - script = _AttemptScript(approach, [ - ("best_effort", 0.40), - ("error", AgentSessionFatalError("3 consecutive agent queries died")), - ]) - approach._solve_attempt = script - with pytest.raises(AgentSessionFatalError): - approach._solve(task, timeout=10) - assert script.calls == 2 # attempt 3 never runs - assert ctx.attempt_start is None - assert ctx.attempt_deadline is None - assert ctx.attempt_index == 0 - - -def test_unexpected_error_without_bank_reraises_after_cleanup(): - """ApproachTimeout (an ApproachFailure SIBLING) propagates, but only. - - after attempt bookkeeping is cleared - stale fields would pollute - later sessions sharing the ToolContext. - """ - approach, task = _make_approach({ - "agent_solve_max_attempts": 3, - "agent_solve_attempt_wall_clock": 2700, - }) - ctx = approach._tool_context - script = _AttemptScript(approach, [("error", ApproachTimeout("slow"))]) - approach._solve_attempt = script - with pytest.raises(ApproachTimeout): - approach._solve(task, timeout=10) - assert ctx.attempt_start is None - assert ctx.attempt_deadline is None - assert ctx.attempt_index == 0 - - -def test_journal_task_context_written_once_even_on_resolve(tmp_path): - """Re-entering _solve for the same task (mid-episode replan) must not - duplicate the goal + init-state entry.""" - approach, task = _make_approach( - { - "agent_solve_max_attempts": 1, - "agent_solve_use_journal": True, - }, - sandbox_dir=str(tmp_path)) - approach._tool_context.test_task_idx = 0 - script = _AttemptScript(approach, [("validated", 0.95), - ("validated", 0.96)]) - approach._solve_attempt = script - approach._solve(task, timeout=10) - approach._solve(task, timeout=10) - content = journal_mod.read_journal(str(tmp_path), - filename=journal_mod.ATTEMPTS_FILENAME) - assert content.count("### cycle 0 task 0 goal + initial state (auto)") == 1 - assert content.index("- goal:") < content.index("- outcome:") - - -def test_test_phase_journal_archived_and_rolled_back(tmp_path): - """Learning content persists across evaluations; each evaluation's own - additions (harness attempt-log entries and agent journal notes) are - archived outside the sandbox, then rolled back so the next evaluation - starts from learning knowledge only (no test-task leaks).""" - sandbox = tmp_path / "sandbox" - log_dir = tmp_path / "run_logs" - approach, task = _make_approach( - { - "agent_solve_max_attempts": 1, - "agent_solve_use_journal": True, - "log_file": str(log_dir), - }, - sandbox_dir=str(sandbox)) - ctx = approach._tool_context - attempts = journal_mod.ATTEMPTS_FILENAME - # A learning-phase note, written by the agent before any evaluation. - sandbox.mkdir(parents=True, exist_ok=True) - (sandbox / journal_mod.JOURNAL_FILENAME).write_text( - "### learn cycle notes\n- learning fact\n", encoding="utf-8") - # First evaluation: one test-task solve writes attempt-log entries - # and the agent adds a note. - approach.begin_test_phase() - ctx.test_task_idx = 0 - script = _AttemptScript(approach, [("validated", 0.95), - ("validated", 0.96)]) - approach._solve_attempt = script - approach._solve(task, timeout=10) - with open(sandbox / journal_mod.JOURNAL_FILENAME, "a", - encoding="utf-8") as f: - f.write("### task 0 attempt 1\n- eval-time note\n") - content = journal_mod.read_journal(str(sandbox), filename=attempts) - assert "### cycle 0 task 0 goal + initial state (auto)" in content - assert "- learning fact" in journal_mod.read_journal(str(sandbox)) - approach.end_test_phase() - # Rolled back: learning content survives, eval additions are gone. - assert "task 0" not in journal_mod.read_journal(str(sandbox), - filename=attempts) - notes = journal_mod.read_journal(str(sandbox)) - assert "- learning fact" in notes - assert "eval-time note" not in notes - # Both files were archived outside the sandbox first, one copy per - # evaluation phase. This first evaluation precedes any online - # learning, so it archives as the initial test. - archived = (log_dir / "attempts_eval_initial.md").read_text() - assert "### cycle 0 task 0 goal + initial state (auto)" in archived - archived_notes = (log_dir / "journal_eval_initial.md").read_text() - assert "- learning fact" in archived_notes - assert "eval-time note" in archived_notes - # Second evaluation on the same task, after a learning phase advanced - # the cycle: the context-entry dedup key was rolled back too, so the - # goal + init entry is re-written (else the attempt records would be - # uninterpretable). - approach._online_learning_cycle = 1 - approach.begin_test_phase() - ctx.test_task_idx = 0 - approach._solve(task, timeout=10) - content = journal_mod.read_journal(str(sandbox), filename=attempts) - assert content.count("### cycle 0 task 0 goal + initial state (auto)") == 1 - approach.end_test_phase() - # The second evaluation ran after cycle 0's learn advanced the - # counter to 1, so it archives under the 0-based cycle it evaluates. - assert sorted(p.name for p in log_dir.glob("attempts_eval*.md")) == [ - "attempts_eval_cycle0.md", "attempts_eval_initial.md" - ] - assert journal_mod.read_raw(str(sandbox)) is not None - assert "task 0" not in journal_mod.read_journal(str(sandbox), - filename=attempts) - - -def test_test_phase_journal_rollback_noop_without_journal(tmp_path): - """With the journal disabled the phase hooks touch nothing.""" - approach, _task = _make_approach({"agent_solve_use_journal": False}, - sandbox_dir=str(tmp_path)) - approach.begin_test_phase() - approach.end_test_phase() - assert journal_mod.read_raw(str(tmp_path)) is None - - -def test_journal_records_best_refused_submission(tmp_path): - """An attempt with no capture journals the best refused submission the - tools stashed, so the plan survives the fresh-context restart.""" - approach, task = _make_approach( - { - "agent_solve_max_attempts": 1, - "agent_solve_use_journal": True, - }, - sandbox_dir=str(tmp_path)) - ctx = approach._tool_context - ctx.test_task_idx = 0 - - def _attempt(_task): - ctx.best_uncaptured_plan_lines = ["Move(block0:block)[0.87]"] - ctx.best_uncaptured_reward = -0.05 - raise ApproachFailure("no capture") - - approach._solve_attempt = _attempt - with pytest.raises(ApproachFailure): - approach._solve(task, timeout=10) - content = journal_mod.read_journal(str(tmp_path), - filename=journal_mod.ATTEMPTS_FILENAME) - assert ("- best refused submission (evaluator reward -0.05, " - "not captured):") in content - assert "Move(block0:block)[0.87]" in content diff --git a/tests/approaches/test_gnn_dynamics_approach.py b/tests/approaches/test_gnn_dynamics_approach.py deleted file mode 100644 index 70d7a00580..0000000000 --- a/tests/approaches/test_gnn_dynamics_approach.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Tests for the GNN dynamics + shooting baseline (paper arm C5).""" -# pylint: disable=protected-access -import numpy as np -import pytest - -from predicators import utils -from predicators.approaches import ApproachFailure, ApproachTimeout, \ - create_approach -from predicators.datasets import create_dataset -from predicators.envs import create_new_env -from predicators.ground_truth_models import get_gt_options -from predicators.settings import CFG -from predicators.structs import InteractionResult - - -def _setup(env_name: str = "cover"): - utils.reset_config({ - "env": env_name, - "num_train_tasks": 3, - "num_test_tasks": 2, - "gnn_num_epochs": 20, - "gnn_use_validation_set": False, - "gnn_do_normalization": True, - "gnn_dynamics_history_len": 2, - "gnn_dynamics_shooting_max_tries": 5, - "gnn_dynamics_max_plan_length": 4, - "explorer": "random_options", - "online_nsrt_learning_requests_per_cycle": 1, - "max_num_steps_interaction_request": 5, - "horizon": 10, - "timeout": 5, - }) - env = create_new_env(env_name) - train_tasks = [t.task for t in env.get_train_tasks()] - options = get_gt_options(env.get_name()) - approach = create_approach("gnn_dynamics_shooting", env.predicates, - options, env.types, env.action_space, - train_tasks) - predicates, _ = utils.parse_config_excluded_predicates(env) - dataset = create_dataset(env, train_tasks, options, predicates) - return env, train_tasks, options, approach, dataset - - -def test_gnn_dynamics_learns_predicts_and_plans(): - """Learn from demos, predict a transition, shoot a plan, save and load.""" - env, train_tasks, options, approach, dataset = _setup() - assert approach.is_learning_based - task = env.get_test_tasks()[0].task - with pytest.raises(ApproachFailure): # nothing learned yet - approach.solve(task, timeout=CFG.timeout) - approach.learn_from_offline_dataset(dataset) - assert approach._gnn is not None - # Option-level examples came out of the demos with history attached. - examples = approach._generate_examples() - assert examples - assert all(len(h) <= 2 for h, *_ in examples) - history, state, option, next_state, num_actions = examples[-1] - pred_state, pred_steps = approach.predict_next_state( - history, state, option) - assert set(pred_state) == set(state) - assert pred_steps >= 1 - assert isinstance(num_actions, int) - for obj in state: - assert pred_state[obj].shape == next_state[obj].shape - # A shooting policy is returned; executing it either reaches the goal - # or fails honestly (the tiny model is not expected to be accurate). - try: - policy = approach.solve(task, timeout=CFG.timeout) - utils.run_policy_with_simulator(policy, - env.simulate, - task.init, - task.goal_holds, - max_num_steps=CFG.horizon, - exceptions_to_break_on={ - utils.OptionExecutionFailure, - ApproachFailure, - }) - except (ApproachFailure, ApproachTimeout): - pass - # Save / load round trip rebuilds the model. - approach2 = create_approach("gnn_dynamics_shooting", env.predicates, - options, env.types, env.action_space, - train_tasks) - approach2.load(online_learning_cycle=None) - assert approach2._gnn is not None - assert approach2._feat_to_index == approach._feat_to_index - s2, _ = approach2.predict_next_state(history, state, option) - for obj in state: - assert np.allclose(s2[obj], pred_state[obj]) - - -def test_gnn_dynamics_online_cycle(): - """Interaction requests come from the configured explorer and the results - extend the data the next model is trained on.""" - env, _, _, approach, dataset = _setup() - approach.learn_from_offline_dataset(dataset) - n_before = len(approach._trajectories) - requests = approach.get_interaction_requests() - assert len(requests) == 1 - request = requests[0] - task = approach._train_tasks[request.train_task_idx] - traj, _ = utils.run_policy(request.act_policy, - env, - "train", - request.train_task_idx, - request.termination_function, - max_num_steps=5, - exceptions_to_break_on={ - utils.RequestActPolicyFailure, - }) - del task - result = InteractionResult(traj.states, traj.actions, - [None] * len(traj.states)) - approach.learn_from_interaction_results([result]) - assert len(approach._trajectories) == n_before + 1 - assert approach._online_learning_cycle == 1 - # The interaction trajectory is not a demo and keeps its task index. - assert not approach._trajectories[-1].is_demo - assert approach._trajectories[-1].train_task_idx == request.train_task_idx - - -def test_gnn_dynamics_no_transitions_keeps_model(): - """With no option-bearing data the previous model is kept.""" - _, _, _, approach, dataset = _setup() - approach.learn_from_offline_dataset(dataset) - gnn = approach._gnn - approach._trajectories = [] - approach._learn_model() - assert approach._gnn is gnn diff --git a/tests/approaches/test_published_fit_reuse.py b/tests/approaches/test_published_fit_reuse.py index 4db849f2c8..148f07f991 100644 --- a/tests/approaches/test_published_fit_reuse.py +++ b/tests/approaches/test_published_fit_reuse.py @@ -79,7 +79,6 @@ def test_reused_physics_fit_restores_the_margin_gate_state( (applying resets them), so the capture gate's margin sweep survives the skip of the harness refit.""" utils.reset_config({ - "agent_sim_learn_oracle_sim_params": False, "agent_explorer_info_seeking": False, }) sim_file = tmp_path / "simulator.py" @@ -94,7 +93,6 @@ def test_reused_physics_fit_restores_the_margin_gate_state( # and make the identity assert below unreachable for mypy. setattr(approach, "_last_fit_result", None) approach._fit_sse = float("inf") - approach._cycle_applied_physical = {} approach._identified_physical_sigma_points = [] approach._rng = np.random.default_rng(0) applied_calls = [] @@ -127,7 +125,6 @@ def _paths() -> Any: assert approach._last_fit_result is fit assert approach._fit_sse == 0.5 assert applied_calls == [{"mu": 0.7}] - assert approach._cycle_applied_physical == {"mu": 0.7} assert approach._identified_physical_sigma_points == sigma @@ -147,7 +144,7 @@ def test_publish_without_a_fit_result_never_deploys(tmp_path: Any) -> None: def test_unfitted_deployment_carries_values_and_clears_evidence( tmp_path: Any, monkeypatch: Any, has_data: bool) -> None: """Edits retain compatible values, initialize new specs, and retire the old - model's posterior without calling any fitting backend.""" + model's posterior.""" utils.reset_config({"agent_sim_learn_param_uncertainty": False}) sim_file = tmp_path / "simulator.py" sim_file.write_text("before", encoding="utf-8") @@ -161,8 +158,6 @@ def test_unfitted_deployment_carries_values_and_clears_evidence( setattr(approach, "_last_fit_result", old_fit) approach._identified_physical_params = {"mu": .7, "removed": 9.} approach._identified_physical_sigma_points = [{"mu": .6}] - approach._cycle_applied_physical = dict( - approach._identified_physical_params) approach._physical_param_specs = [ParamSpec("mu", .5, 0., 1.)] specs = [ ParamSpec("k", 1., 0., 2.), @@ -179,20 +174,11 @@ def apply(params): approach._identified_physical_params = dict(params) approach._identified_physical_sigma_points = [] - def forbidden(*_args, **_kwargs): - pytest.fail("Deployment must not fit") - monkeypatch.setattr(approach, "_apply_identified_physical_params", apply) - monkeypatch.setattr(approach, "_fit_parameters_joint_rollout", forbidden) - monkeypatch.setattr(approach, "_fit_parameters_recurrent", forbidden) - monkeypatch.setattr( - "predicators.approaches.agent_sim_learning_approach" - ".fit_rule_parameters", forbidden) triples: Any = [(None, None, None)] if has_data else [] approach._fit_params_after_synthesis([], specs, triples, {}) assert approach._fitted_params == {"k": 1.5, "bounded": 2., "new": .25} assert applied == [{"mu": .7}] - assert approach._cycle_applied_physical == {"mu": .7} assert not approach._identified_physical_sigma_points assert approach._last_fit_result is None assert approach._fit_sse == float("inf") diff --git a/tests/approaches/test_sampler_learning_mixin.py b/tests/approaches/test_sampler_learning_mixin.py deleted file mode 100644 index 87c1d8829e..0000000000 --- a/tests/approaches/test_sampler_learning_mixin.py +++ /dev/null @@ -1,174 +0,0 @@ -"""Tests for SamplerLearningMixin's loader and oracle-install logic. - -Covers ``_load_samplers_from_module_file`` (missing file, exec error, -non-dict, bad entries, happy path) and -``_maybe_install_oracle_samplers`` (GT install, fallback to synthesis, -disabled no-op) on a minimal host. -""" - -from typing import Any, Dict, Set - -import numpy as np -from gym.spaces import Box - -from predicators import utils -from predicators.approaches import sampler_learning_mixin -from predicators.approaches.sampler_learning_mixin import SamplerLearningMixin -from predicators.structs import Action, GroundAtom, Object, \ - ParameterizedOption, Predicate, State, Task, Type - -_block_type = Type("block", ["x"]) -_block = Object("block0", _block_type) - -_Reached = Predicate("Reached", [_block_type], lambda s, o: True) - -_Move = ParameterizedOption( - "Move", - types=[_block_type], - params_space=Box(low=np.array([0.0], dtype=np.float32), - high=np.array([1.0], dtype=np.float32)), - policy=lambda _s, _m, _o, _p: Action(np.zeros(1, dtype=np.float32)), - initiable=lambda _s, _m, _o, _p: True, - terminal=lambda _s, _m, _o, _p: False, -) - - -class _Host(SamplerLearningMixin): # pylint: disable=abstract-method - """Minimal host supplying the mixin's contract surface.""" - - def __init__(self): - init = State({_block: np.array([0.0], dtype=np.float32)}) - self._types = {_block_type} - self._train_tasks = [Task(init, {GroundAtom(_Reached, [_block])})] - self._fitted_params: Dict[str, float] = {} - self._synthesized_samplers: Dict[str, Any] = {} - self._init_sampler_learning_state() - - def _get_all_predicates(self) -> Set[Predicate]: - return {_Reached} - - def _get_all_options(self) -> Set[ParameterizedOption]: - return {_Move} - - def _learning_cycle_index(self) -> int: - return 1 - - -def _host(**config): - utils.reset_config({"seed": 0, **config}) - return _Host() - - -# --------------------------------------------------------------------------- # -# _load_samplers_from_module_file. -# --------------------------------------------------------------------------- # - - -def test_load_samplers_missing_file_returns_empty(tmp_path): - """A missing samplers.py loads as the empty dict (samplers optional).""" - host = _host() - assert not host._load_samplers_from_module_file( # pylint: disable=protected-access - str(tmp_path / "samplers.py")) - - -def test_load_samplers_exec_error_returns_empty(tmp_path): - """A file that raises at exec time loads as the empty dict.""" - path = tmp_path / "samplers.py" - path.write_text("raise RuntimeError('boom')\n", encoding="utf-8") - host = _host() - assert not host._load_samplers_from_module_file( # pylint: disable=protected-access - str(path)) - - -def test_load_samplers_non_dict_returns_empty(tmp_path): - """LEARNED_SAMPLERS bound to a non-dict loads as the empty dict.""" - path = tmp_path / "samplers.py" - path.write_text("LEARNED_SAMPLERS = [1, 2]\n", encoding="utf-8") - host = _host() - assert not host._load_samplers_from_module_file( # pylint: disable=protected-access - str(path)) - - -def test_load_samplers_skips_unknown_and_non_callable_entries(tmp_path): - """Unknown option names and non-callables are dropped, the rest kept.""" - path = tmp_path / "samplers.py" - path.write_text("""\ -def _fn(state, subgoal_atoms, rng, objects): - del state, subgoal_atoms, objects - return np.array([0.5], dtype=np.float32) - -LEARNED_SAMPLERS = {"Move": _fn, "Teleport": _fn, "Reached": 7} -""", - encoding="utf-8") - host = _host() - loaded = host._load_samplers_from_module_file(str(path)) # pylint: disable=protected-access - assert set(loaded) == {"Move"} - - -def test_load_samplers_happy_path(tmp_path): - """A valid file loads a callable that draws correctly shaped params.""" - path = tmp_path / "samplers.py" - path.write_text("""\ -def _fn(state, subgoal_atoms, rng, objects): - del state, subgoal_atoms, objects - return np.array([0.25 + 0.01 * rng.random()], dtype=np.float32) - -LEARNED_SAMPLERS = {"Move": _fn} -""", - encoding="utf-8") - host = _host() - loaded = host._load_samplers_from_module_file(str(path)) # pylint: disable=protected-access - assert set(loaded) == {"Move"} - draw = loaded["Move"]( - host._train_tasks[0].init, # pylint: disable=protected-access - set(), - np.random.default_rng(0), - [_block]) - assert np.asarray(draw).shape == (1, ) - - -# --------------------------------------------------------------------------- # -# _maybe_install_oracle_samplers. -# --------------------------------------------------------------------------- # - - -def _gt_sampler(state, subgoal_atoms, rng, objects): - del state, subgoal_atoms, rng, objects - return np.array([0.5], dtype=np.float32) - - -def test_oracle_samplers_installed_when_available(monkeypatch): - """With oracle_samplers on and GT available: install, skip synthesis.""" - monkeypatch.setattr(sampler_learning_mixin, "get_gt_samplers", - lambda _env: {"Move": _gt_sampler}) - host = _host(agent_sim_learn_parameterized_samplers=True, - agent_sim_learn_oracle_samplers=True) - host._maybe_install_oracle_samplers() # pylint: disable=protected-access - assert host._synthesized_samplers == {"Move": _gt_sampler} # pylint: disable=protected-access - assert host._current_samplers_version == "oracle" # pylint: disable=protected-access - assert not host._do_synthesize_samplers # pylint: disable=protected-access - - -def test_oracle_samplers_fall_back_to_synthesis_when_none(monkeypatch): - """With oracle_samplers on but no GT for the env: synthesize instead.""" - monkeypatch.setattr(sampler_learning_mixin, "get_gt_samplers", - lambda _env: {}) - host = _host(agent_sim_learn_parameterized_samplers=True, - agent_sim_learn_oracle_samplers=True) - host._maybe_install_oracle_samplers() # pylint: disable=protected-access - assert not host._synthesized_samplers # pylint: disable=protected-access - assert host._do_synthesize_samplers # pylint: disable=protected-access - - -def test_samplers_disabled_no_synthesis_no_install(monkeypatch): - """With the master gate off nothing is installed or synthesized.""" - - def _boom(_env): - raise AssertionError("get_gt_samplers called with samplers disabled") - - monkeypatch.setattr(sampler_learning_mixin, "get_gt_samplers", _boom) - host = _host(agent_sim_learn_parameterized_samplers=False, - agent_sim_learn_oracle_samplers=False) - host._maybe_install_oracle_samplers() # pylint: disable=protected-access - assert not host._synthesized_samplers # pylint: disable=protected-access - assert not host._do_synthesize_samplers # pylint: disable=protected-access diff --git a/tests/approaches/test_sim_learning_info_seeking.py b/tests/approaches/test_sim_learning_info_seeking.py index 27ea816d0c..cc4dc2ee3c 100644 --- a/tests/approaches/test_sim_learning_info_seeking.py +++ b/tests/approaches/test_sim_learning_info_seeking.py @@ -156,37 +156,6 @@ def test_rebuild_param_ensemble_respects_flag(): assert approach._param_ensemble[0] == {"a": 1.0} # member 0 is anchor -def test_rebuild_param_ensemble_empty_under_oracle_params(): - """Oracle params carry no uncertainty, so no ensemble is built. - - Without this the uniform-jitter fallback would hand the capture gate - members that no plan can satisfy (a zero rate, a rewired lamp), and - the gate would refuse every plan on a model that is exactly right. - """ - from predicators.code_sim_learning.fit_space import ParamSpec - approach = object.__new__(AgentSimLearningApproach) - approach._fitted_params = {"a": 1.0} - approach._param_specs = [ParamSpec("a", 1.0, lo=0.0, hi=2.0)] - approach._param_ensemble = [{"a": 1.0}, {"a": 2.0}] - approach._last_fit_result = None - approach._rng = np.random.default_rng(0) - utils.reset_config({ - "agent_plan_validation_rule_param_margin": True, - "agent_explorer_info_ensemble_size": 5, - "agent_sim_learn_oracle_sim_params": True, - }) - approach._rebuild_param_ensemble() - assert approach._param_ensemble == [] - - utils.reset_config({ - "agent_plan_validation_rule_param_margin": True, - "agent_explorer_info_ensemble_size": 5, - "agent_sim_learn_oracle_sim_params": False, - }) - approach._rebuild_param_ensemble() - assert len(approach._param_ensemble) == 5 - - def _selector_approach(fit_result): from predicators.code_sim_learning.fit_space import ParamSpec approach = object.__new__(AgentSimLearningApproach) @@ -261,14 +230,9 @@ def test_select_ensemble_uniform_when_calibration_disabled(): assert method == "uniform-perturb" -def test_fit_params_no_data_seeds_declared_inits(monkeypatch): - """With no transitions, params seed from inits and no fit runs. - - This is the oracle-sim-program no-demos path: every demo failed, so - ``_learn_simulator`` reaches the fit with empty - ``base_pred_triples`` and must fall back to the declared init values - instead of fitting. - """ +def test_fit_params_no_data_seeds_declared_inits(): + """With no transitions and no published fit, the deployed params are the + declared init values.""" from predicators.code_sim_learning.fit_space import ParamSpec approach = object.__new__(AgentSimLearningApproach) @@ -280,16 +244,7 @@ def test_fit_params_no_data_seeds_declared_inits(monkeypatch): approach._last_fit_result = None approach._fit_sse = 0.0 approach._rng = np.random.default_rng(0) - - def _fail_fit(*args, **kwargs): - del args, kwargs - raise AssertionError("fit must not run with no data") - - monkeypatch.setattr( - "predicators.approaches.agent_sim_learning_approach" - ".fit_rule_parameters", _fail_fit) utils.reset_config({ - "agent_sim_learn_oracle_sim_params": False, "agent_explorer_info_seeking": False, }) specs = [ParamSpec("a", 1.5, lo=0.0, hi=5.0)] diff --git a/tests/code_sim_learning/test_param_fitting.py b/tests/code_sim_learning/test_param_fitting.py index 236b94e4ea..baf24ff616 100644 --- a/tests/code_sim_learning/test_param_fitting.py +++ b/tests/code_sim_learning/test_param_fitting.py @@ -13,7 +13,7 @@ import predicators.approaches # noqa: F401 # pylint: disable=unused-import from predicators import utils -from predicators.approaches.agent_model_based_approach import _SketchStep +from predicators.agent_sdk.sketch_types import SketchStep as _SketchStep from predicators.code_sim_learning.fit_space import ParamSpec from predicators.code_sim_learning.fitting import compute_sse, fit_params from predicators.envs import create_new_env diff --git a/tests/execution_monitoring/test_execution_monitoring.py b/tests/execution_monitoring/test_execution_monitoring.py index 422f91b5f1..6c7c5acfe2 100644 --- a/tests/execution_monitoring/test_execution_monitoring.py +++ b/tests/execution_monitoring/test_execution_monitoring.py @@ -1,20 +1,14 @@ """Tests for execution monitors.""" -import numpy as np import pytest -from gym.spaces import Box from predicators.execution_monitoring import create_execution_monitor from predicators.execution_monitoring.expected_atoms_monitor import \ ExpectedAtomsExecutionMonitor from predicators.execution_monitoring.mpc_execution_monitor import \ MpcExecutionMonitor -from predicators.execution_monitoring.subgoal_annotations_monitor import \ - SubgoalAnnotationsExecutionMonitor, SubgoalExecutionStatus from predicators.execution_monitoring.trivial_execution_monitor import \ TrivialExecutionMonitor -from predicators.structs import Action, GroundAtom, Object, \ - ParameterizedOption, Predicate, State, Type def test_create_execution_monitor(): @@ -28,131 +22,6 @@ def test_create_execution_monitor(): exec_monitor = create_execution_monitor("expected_atoms") assert isinstance(exec_monitor, ExpectedAtomsExecutionMonitor) - exec_monitor = create_execution_monitor("subgoal_annotations") - assert isinstance(exec_monitor, SubgoalAnnotationsExecutionMonitor) - with pytest.raises(NotImplementedError) as e: create_execution_monitor("not a real monitor") assert "Unrecognized execution monitor" in str(e) - - -class _FakeSketchStep: - """Duck-typed sketch step (see agent_sdk.bilevel_sketch.SketchStep).""" - - def __init__(self, option, subgoal_atoms, subgoal_neg_atoms=None): - self.option = option - self.subgoal_atoms = subgoal_atoms - self.subgoal_neg_atoms = subgoal_neg_atoms - - -def test_subgoal_annotations_monitor(): - """Unit tests for SubgoalAnnotationsExecutionMonitor.step().""" - block_type = Type("block", ["held"]) - block = Object("block0", block_type) - held = Predicate("Held", [block_type], - lambda s, o: s.get(o[0], "held") > 0.5) - state_held = State({block: np.array([1.0], dtype=np.float32)}) - state_free = State({block: np.array([0.0], dtype=np.float32)}) - - def _make_option(terminal): - param_opt = ParameterizedOption( - "Pick", - types=[block_type], - params_space=Box(low=np.zeros(1, dtype=np.float32), - high=np.ones(1, dtype=np.float32)), - policy=lambda s, m, o, p: Action(np.zeros(1, dtype=np.float32)), - initiable=lambda s, m, o, p: True, - terminal=lambda s, m, o, p: terminal, - ) - return param_opt, param_opt.ground([block], - np.zeros(1, dtype=np.float32)) - - done_parent, done_option = _make_option(True) - _, running_option = _make_option(False) - held_atom = GroundAtom(held, [block]) - - monitor = create_execution_monitor("subgoal_annotations") - - # No approach info (e.g. exploration): never replan. - assert not monitor.step(state_free) - - # Info of an unexpected shape (another approach's export): ignore. - monitor.update_approach_info([{"something": "else"}]) - assert not monitor.step(state_free) - - def _status(option, steps_initiated, pos=None, neg=None): - step = _FakeSketchStep(done_parent, pos, neg) - return SubgoalExecutionStatus(sketch=[step], - steps_initiated=steps_initiated, - current_option=option) - - # No option initiated yet (fresh policy right after a replan). - monitor.update_approach_info([_status(None, 0, {held_atom})]) - assert not monitor.step(state_free) - - # Mid-option: the current option has not terminated. - monitor.update_approach_info([_status(running_option, 1, {held_atom})]) - assert not monitor.step(state_free) - - # Boundary, annotation holds: no replan. - monitor.update_approach_info([_status(done_option, 1, {held_atom})]) - assert not monitor.step(state_held) - - # Boundary, unannotated step: nothing to check. - monitor.update_approach_info([_status(done_option, 1, None)]) - assert not monitor.step(state_free) - - # Boundary, positive atom unsatisfied: replan. - monitor.update_approach_info([_status(done_option, 1, {held_atom})]) - assert monitor.step(state_free) - - # Boundary, negative atom violated: replan. - monitor.update_approach_info([_status(done_option, 1, None, {held_atom})]) - assert monitor.step(state_held) - - -def test_subgoal_annotations_monitor_skips_unverifiable_atoms(): - """A classifier that cannot evaluate on a bare observation (e.g. it indexes - a latent that real env states never carry) is skipped with a warning - instead of crashing the episode or counting as divergence.""" - block_type = Type("block", ["held"]) - block = Object("block0", block_type) - - def _latent_only(s, o, latent=None): - del s, o - return bool(latent["_bonds"]) # TypeError when latent is None - - latent_pred = Predicate("LatentBonded", [block_type], _latent_only) - latent_atom = GroundAtom(latent_pred, [block]) - held = Predicate("Held", [block_type], - lambda s, o: s.get(o[0], "held") > 0.5) - held_atom = GroundAtom(held, [block]) - state = State({block: np.array([0.0], dtype=np.float32)}) # no latent - - param_opt = ParameterizedOption( - "Pick", - types=[block_type], - params_space=Box(low=np.zeros(1, dtype=np.float32), - high=np.ones(1, dtype=np.float32)), - policy=lambda s, m, o, p: Action(np.zeros(1, dtype=np.float32)), - initiable=lambda s, m, o, p: True, - terminal=lambda s, m, o, p: True, - ) - done_option = param_opt.ground([block], np.zeros(1, dtype=np.float32)) - monitor = create_execution_monitor("subgoal_annotations") - - def _status(pos=None, neg=None): - step = _FakeSketchStep(param_opt, pos, neg) - return SubgoalExecutionStatus(sketch=[step], - steps_initiated=1, - current_option=done_option) - - # The unverifiable atom alone: skipped, no divergence, no crash. - monitor.update_approach_info([_status(pos={latent_atom})]) - assert not monitor.step(state) - # Same in the negative polarity. - monitor.update_approach_info([_status(neg={latent_atom})]) - assert not monitor.step(state) - # A genuinely failed observable atom alongside it still fires. - monitor.update_approach_info([_status(pos={latent_atom, held_atom})]) - assert monitor.step(state) diff --git a/tests/explorers/test_agent_model_based_explorer.py b/tests/explorers/test_agent_model_based_explorer.py deleted file mode 100644 index 3af97c9330..0000000000 --- a/tests/explorers/test_agent_model_based_explorer.py +++ /dev/null @@ -1,491 +0,0 @@ -"""Tests for AgentModelBasedExplorer.""" -# pylint: disable=protected-access - -from unittest.mock import AsyncMock, MagicMock - -import numpy as np -import pytest -from gym.spaces import Box - -from predicators import utils -from predicators.agent_sdk.sketch_types import SketchStep -from predicators.agent_sdk.tools import ToolContext -from predicators.explorers import create_explorer -from predicators.explorers.agent_model_based_explorer import \ - AgentModelBasedExplorer -from predicators.explorers.base_explorer import BaseExplorer -from predicators.structs import Action, GroundAtom, Object, \ - ParameterizedOption, Predicate, State, Task, Type - -# --------------------------------------------------------------------------- -# Fixtures (parallel the bilevel approach tests) -# --------------------------------------------------------------------------- - -_block_type = Type("block", ["x", "y", "held"]) -_robot_type = Type("robot", ["x", "y"]) - -_block0 = Object("block0", _block_type) -_block1 = Object("block1", _block_type) -_robot = Object("robot0", _robot_type) - -_Holding = Predicate("Holding", [_block_type], - lambda s, o: s.get(o[0], "held") > 0.5) -_On = Predicate("On", [_block_type, _block_type], - lambda s, o: abs(s.get(o[0], "x") - s.get(o[1], "x")) < 0.1) -_HandEmpty = Predicate("HandEmpty", [_robot_type], lambda s, o: True) - -_ALL_PREDICATES = {_Holding, _On, _HandEmpty} -_ALL_TYPES = {_block_type, _robot_type} - - -def _noop_policy(_s, _m, _o, _p): - return Action(np.zeros(1, dtype=np.float32)) - - -def _always_true(_s, _m, _o, _p): - return True - - -def _always_false(_s, _m, _o, _p): - return False - - -_Pick = ParameterizedOption( - "Pick", - types=[_block_type], - params_space=Box(low=np.array([0.0], dtype=np.float32), - high=np.array([1.0], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_true, - terminal=_always_false, -) - -_Place = ParameterizedOption( - "Place", - types=[_block_type, _block_type], - params_space=Box(low=np.array([0.0, 0.0], dtype=np.float32), - high=np.array([1.0, 1.0], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_true, - terminal=_always_false, -) - -_Wait = ParameterizedOption( - "Wait", - types=[_robot_type], - params_space=Box(low=np.array([], dtype=np.float32), - high=np.array([], dtype=np.float32)), - policy=_noop_policy, - initiable=_always_true, - terminal=_always_false, -) - -_ALL_OPTIONS = {_Pick, _Place, _Wait} - - -def _make_state(overrides=None): - data = { - _block0: np.array([0.1, 0.2, 0.0], dtype=np.float32), - _block1: np.array([0.5, 0.6, 0.0], dtype=np.float32), - _robot: np.array([0.0, 0.0], dtype=np.float32), - } - if overrides: - for obj, vals in overrides.items(): - data[obj] = np.array(vals, dtype=np.float32) - return State(data) - - -def _make_task(): - state = _make_state() - goal = {GroundAtom(_On, [_block0, _block1])} - return Task(state, goal) - - -def _assistant_response(text: str): - return [{ - "type": "assistant", - "content": [{ - "type": "text", - "text": text - }], - }] - - -def _make_explorer(option_model, query_impl): - """Build an AgentModelBasedExplorer with stubbed session + tool_context.""" - tool_context = ToolContext( - types=_ALL_TYPES, - predicates=_ALL_PREDICATES, - options=_ALL_OPTIONS, - train_tasks=[_make_task()], - option_model=option_model, - ) - agent_session = MagicMock() - agent_session.query = query_impl - agent_session.tool_names = None - explorer = AgentModelBasedExplorer( - predicates=_ALL_PREDICATES, - options=_ALL_OPTIONS, - types=_ALL_TYPES, - action_space=Box(low=-1, high=1, shape=(1, )), - train_tasks=[_make_task()], - max_steps_before_termination=50, - tool_context=tool_context, - agent_session=agent_session, - ) - return explorer, tool_context - - -def _reset_config(**overrides): - base = { - "env": "cover", - "approach": "agent_model_based", - "num_train_tasks": 1, - "num_test_tasks": 1, - "seed": 42, - "agent_bilevel_max_samples_per_step": 5, - "agent_bilevel_check_subgoals": True, - "agent_bilevel_log_state": False, - "agent_explorer_fallback_to_random": True, - "agent_sdk_max_trajectories_in_context": 5, - } - base.update(overrides) - utils.reset_config(base) - - -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - - -def test_factory_registration(): - """AgentModelBasedExplorer is reachable through create_explorer, under its - own name and under the deprecated agent_bilevel alias.""" - _reset_config() - tool_context = ToolContext( - types=_ALL_TYPES, - predicates=_ALL_PREDICATES, - options=_ALL_OPTIONS, - train_tasks=[_make_task()], - option_model=MagicMock(), - ) - agent_session = MagicMock() - explorer = create_explorer( - "agent_model_based", - _ALL_PREDICATES, - _ALL_OPTIONS, - _ALL_TYPES, - Box(low=-1, high=1, shape=(1, )), - [_make_task()], - tool_context=tool_context, - agent_session=agent_session, - ) - assert isinstance(explorer, BaseExplorer) - assert isinstance(explorer, AgentModelBasedExplorer) - legacy = create_explorer( - "agent_bilevel", - _ALL_PREDICATES, - _ALL_OPTIONS, - _ALL_TYPES, - Box(low=-1, high=1, shape=(1, )), - [_make_task()], - tool_context=tool_context, - agent_session=agent_session, - ) - assert isinstance(legacy, AgentModelBasedExplorer) - - -def test_happy_path_returns_policy_and_stashes_subgoals(): - """Canned sketch → refined plan → policy and stashed subgoals.""" - _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) - - plan_text = ("Pick(block0:block)\n" - "Place(block0:block, block1:block) -> " - "{On(block0:block, block1:block)}\n") - query = AsyncMock(return_value=_assistant_response(plan_text)) - - explorer, tool_context = _make_explorer(option_model, query) - policy, term_fn = explorer._get_exploration_strategy(0, timeout=5) - - assert callable(policy) - assert term_fn(_make_state()) is False - assert tool_context.last_sketch_subgoals is not None - assert len(tool_context.last_sketch_subgoals) == 2 - # Second step's positive subgoal should be {On(block0, block1)}. - pos2, _neg2 = tool_context.last_sketch_subgoals[1] - assert pos2 == {GroundAtom(_On, [_block0, _block1])} - assert tool_context.last_sketch_options == [ - ("Pick", ["block0"]), - ("Place", ["block0", "block1"]), - ] - assert query.await_count == 1 - - -def test_wait_memory_injection_on_grounding(): - """A Wait step's annotated subgoal rides on the grounded option as - ``wait_target_atoms`` so WaitOption terminates on the intended atoms.""" - _reset_config() - explorer, _ = _make_explorer(MagicMock(), None) - step = SketchStep(option=_Wait, - objects=[_robot], - subgoal_atoms={GroundAtom(_On, [_block0, _block1])}) - plan = explorer._ground_sketch_verbatim([step]) - assert len(plan) == 1 and plan[0].name == "Wait" - assert plan[0].memory["wait_target_atoms"] == { - GroundAtom(_On, [_block0, _block1]) - } - - -def test_sketch_executes_verbatim_without_belief_refinement(): - """The agent's explicit parameters execute exactly as written: the belief - model is never rolled, the verdict is not-certified, and the cycle record - shows the executed values.""" - _reset_config(agent_bilevel_use_llm_initial_params=True) - option_model = MagicMock() - plan_text = ("```\nPick(block0:block)[0.42] -> {Holding(block0:block)}\n" - "Place(block0:block, block1:block)[0.11, 0.22] -> " - "{On(block0:block, block1:block)}\n```") - query = AsyncMock(return_value=_assistant_response(plan_text)) - explorer, tool_context = _make_explorer(option_model, query) - policy, term_fn = explorer._get_exploration_strategy(0, timeout=5) - assert callable(policy) and term_fn(_make_state()) is False - assert not option_model.get_next_state_and_num_actions.called - assert tool_context.last_mental_model_solved is False - record = tool_context.cycle_scheduled_plans[-1] - assert "Pick(block0)[0.4200]" in record - assert "Place(block0, block1)[0.1100, 0.2200]" in record - assert "-> {On(block0:block, block1:block)}" in record - assert "without belief-model certification" in record - assert tool_context.last_sketch_options == [ - ("Pick", ["block0"]), - ("Place", ["block0", "block1"]), - ] - - -def test_missing_params_get_one_draw_from_the_box(): - """A step the agent left without parameters is grounded on one uniform. - - draw from the option's box - no search, and no crash. - """ - _reset_config(agent_bilevel_use_llm_initial_params=True) - explorer, _ = _make_explorer(MagicMock(), None) - steps = [ - SketchStep(option=_Pick, objects=[_block0], subgoal_atoms=None), - SketchStep(option=_Place, - objects=[_block0, _block1], - subgoal_atoms=None, - initial_params=np.array([0.5], dtype=np.float32)), - ] - plan = explorer._ground_sketch_verbatim(steps) - assert plan[0].params.shape == (1, ) and 0.0 <= plan[0].params[0] <= 1.0 - # Wrong arity counts as missing. - assert plan[1].params.shape == (2, ) - - -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 submit_plan but ends in prose: - - explorer recovers the captured plan instead of falling back to - random and executes it at the captured continuous params. - """ - _reset_config() - option_model = MagicMock() - 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 - # The captured params execute verbatim; the belief is not re-rolled. - assert not option_model.get_next_state_and_num_actions.called - record = tool_context.cycle_scheduled_plans[-1] - assert "Pick(block0)[0.4200]" in record - assert "Place(block0, block1)[0.1100, 0.2200]" in record - - -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) - - option_model = MagicMock() - - async def failing_query(_msg): - raise RuntimeError("boom") - - explorer, _ = _make_explorer(option_model, failing_query) - policy, term_fn = explorer._get_exploration_strategy(0, timeout=5) - assert callable(policy) - assert term_fn(_make_state()) is False - - -def test_fallback_disabled_raises(): - """Agent raises → RequestActPolicyFailure when fallback flag off.""" - _reset_config(agent_explorer_fallback_to_random=False) - - option_model = MagicMock() - - async def failing_query(_msg): - raise RuntimeError("boom") - - explorer, _ = _make_explorer(option_model, failing_query) - with pytest.raises(utils.RequestActPolicyFailure): - explorer._get_exploration_strategy(0, timeout=5) - - -def test_experiment_guidance_gated_by_info_seeking(): - """Experiment guidance appears iff info-seeking is on.""" - _reset_config(agent_explorer_info_seeking=True) - explorer, _ = _make_explorer(MagicMock(), MagicMock()) - guidance = explorer._build_experiment_guidance() # pylint: disable=protected-access - assert "sim.suggest_probes" in guidance - # Off => section absent entirely. - _reset_config(agent_explorer_info_seeking=False) - assert explorer._build_experiment_guidance() == "" # pylint: disable=protected-access - - -def test_experiment_guidance_injects_open_questions_ledger(tmp_path): - """The learn phase's open_questions.md ledger reaches the explore query - verbatim, independent of the info-seeking flag, and an oversized ledger - keeps its head (the ranking's top).""" - _reset_config(agent_explorer_info_seeking=False) - explorer, tool_context = _make_explorer(MagicMock(), MagicMock()) - # No sandbox / no file => no section (and no crash). - assert explorer._build_experiment_guidance() == "" # pylint: disable=protected-access - tool_context.sandbox_dir = str(tmp_path) - assert explorer._build_experiment_guidance() == "" # pylint: disable=protected-access - ledger = ("1. Bond window: place pairs at spacings 0.100/0.104/" - "0.110/0.114 and record which bond.") - (tmp_path / "open_questions.md").write_text(ledger, encoding="utf-8") - guidance = explorer._build_experiment_guidance() # pylint: disable=protected-access - assert ledger in guidance - assert "OPEN QUESTIONS" in guidance - assert "The TOP entry is mandatory" in guidance - # Info-seeking on: both the ledger and the boundary-probing note. - _reset_config(agent_explorer_info_seeking=True) - guidance = explorer._build_experiment_guidance() # pylint: disable=protected-access - assert ledger in guidance - assert "sim.suggest_probes" in guidance - # Oversized ledger: head survives, truncation is announced. - head = "TOP-RANKED ENTRY" - (tmp_path / "open_questions.md").write_text(head + "x" * 10000, - encoding="utf-8") - guidance = explorer._build_experiment_guidance() # pylint: disable=protected-access - assert head in guidance - assert "ledger truncated" in guidance - - -def _make_certified_capture(pick_params, place_params): - """A capture as the belief's validation gate leaves it: goal reached.""" - grounded_plan, captured_sketch = _make_captured(pick_params, place_params) - return grounded_plan, captured_sketch - - -def test_certified_capture_executes_verbatim_and_next_request_queries(): - """A plan the session validated through the capture gate (reached_goal - True) is executed verbatim as a solve attempt with a True mental-model - verdict; the cycle's next request on the task queries the agent again - (asking for a different certified plan) rather than replaying it.""" - _reset_config(agent_explorer_info_seeking=True) - option_model = MagicMock() - option_model.get_next_state_and_num_actions.return_value = (_make_state( - {_block0: [0.5, 0.6, 0.0]}), 3) - pick_params, place_params = [0.42], [0.11, 0.22] - grounded_plan, captured_sketch = _make_certified_capture( - pick_params, place_params) - explorer, tool_context = _make_explorer(option_model, None) - tool_context.atom_disagreement_fn = lambda _s, _atoms: 1.0 - queries = [] - - async def query_impl(msg, **_kw): - queries.append(msg) - tool_context.solved_plan = grounded_plan - tool_context.solved_sketch = captured_sketch - tool_context.solved_plan_reached_goal = True - tool_context.solved_plan_validation_summary = "5/5 rollouts ok" - return _assistant_response("Validated 5/5; submitting.") - - explorer._agent_session.query = query_impl - policy, term_fn = explorer._get_exploration_strategy(0, timeout=5) - assert callable(policy) and term_fn(_make_state()) is False - # Verbatim: no refinement rollouts, verdict True, capture consumed. - assert not option_model.get_next_state_and_num_actions.called - assert tool_context.last_mental_model_solved is True - assert tool_context.solved_plan is None - assert "belief-certified" in tool_context.cycle_scheduled_plans[-1] - assert "replayed" not in tool_context.cycle_scheduled_plans[-1] - assert tool_context.last_sketch_options == [("Pick", ["block0"]), - ("Place", ["block0", - "block1"])] - # The policy runs the captured options with their captured params. - act = policy(_make_state()) - assert isinstance(act, Action) - # Second request of the cycle on the same task: a new query that - # shows the certified plan as already scheduled. - tool_context.last_mental_model_solved = None - policy2, _ = explorer._get_exploration_strategy(0, timeout=5) - assert callable(policy2) - assert len(queries) == 2 - assert "belief-certified" in queries[1] - assert "STRUCTURALLY different" in queries[1] - assert tool_context.last_mental_model_solved is True - - -def test_uncertified_capture_executes_its_plan_verbatim(): - """A capture whose gate verdict is not True (best-effort, flaky) is not - certified: it executes at its captured params as an experiment, with a - False mental-model verdict.""" - _reset_config() - option_model = MagicMock() - grounded_plan, captured_sketch = _make_captured([0.42], [0.11, 0.22]) - explorer, tool_context = _make_explorer(option_model, None) - - async def query_impl(_msg, **_kw): - tool_context.solved_plan = grounded_plan - tool_context.solved_sketch = captured_sketch - tool_context.solved_plan_reached_goal = False - return _assistant_response("Best effort only, no sketch block.") - - explorer._agent_session.query = query_impl - policy, _ = explorer._get_exploration_strategy(0, timeout=5) - assert callable(policy) - assert not option_model.get_next_state_and_num_actions.called - assert "belief-certified" not in tool_context.cycle_scheduled_plans[-1] - assert tool_context.last_mental_model_solved is False diff --git a/tests/test_agent_harness_fixes.py b/tests/test_agent_harness_fixes.py index 5d66f73394..13be018edb 100644 --- a/tests/test_agent_harness_fixes.py +++ b/tests/test_agent_harness_fixes.py @@ -29,20 +29,12 @@ def test_max_session_log_number(tmp_path: Path) -> None: assert max_session_log_number(str(tmp_path)) == 7 -def test_real_episode_step_budget_is_phase_aware() -> None: - """Explore episodes are capped by the interaction-request cap too.""" - utils.reset_config({ - "horizon": 3000, - "max_num_steps_interaction_request": 1000 - }) - assert utils.real_episode_step_budget("explore") == 1000 +def test_real_episode_step_budget_is_the_horizon() -> None: + """Every agent session phase gets the horizon.""" + utils.reset_config({"horizon": 3000}) assert utils.real_episode_step_budget("solve") == 3000 + assert utils.real_episode_step_budget("synthesis") == 3000 assert utils.real_episode_step_budget(None) == 3000 - utils.reset_config({ - "horizon": 3000, - "max_num_steps_interaction_request": 5000 - }) - assert utils.real_episode_step_budget("explore") == 3000 _block_type = Type("block", ["x"]) diff --git a/tests/test_main.py b/tests/test_main.py index 65f2587eb6..fda9517ac9 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -13,8 +13,8 @@ from predicators import utils from predicators.approaches import ApproachFailure, ApproachTimeout, \ BaseApproach, create_approach -from predicators.approaches.agent_model_free_approach import \ - AgentModelFreeApproach +from predicators.approaches.pp_online_process_learning_approach import \ + OnlineProcessLearningAndPlanningApproach from predicators.cogman import CogMan from predicators.envs.cover import CoverEnv from predicators.execution_monitoring import create_execution_monitor @@ -465,9 +465,9 @@ def test_stash_resume_restores_request_bookkeeping(): get_interaction_requests, so the result->train-task pairing that learn_from_interaction_results needs must come from restore_interaction_requests (run_20260828_173451 asserted on it).""" - # The model-free family records the pairing in get_interaction_requests - # and asserts on it in learn_from_interaction_results. - approach = object.__new__(AgentModelFreeApproach) + # The online process learner records the pairing in + # get_interaction_requests and reads it in learn_from_interaction_results. + approach = object.__new__(OnlineProcessLearningAndPlanningApproach) approach._requests_train_task_idxs = None # pylint: disable=protected-access approach.restore_interaction_requests([0, 0]) assert approach._requests_train_task_idxs == [0, 0] # pylint: disable=protected-access