diff --git a/contrib/CODEOWNERS b/contrib/CODEOWNERS index fefb88b02..ab5f68ec7 100644 --- a/contrib/CODEOWNERS +++ b/contrib/CODEOWNERS @@ -2,3 +2,7 @@ # Recipes recipes/search_r1 @SiyunZhao @JiahangXu +recipes/shaper @Control-derek + +# Runtime extensions +agentlightning/contrib/shaper @Control-derek diff --git a/contrib/agentlightning/contrib/shaper/README.md b/contrib/agentlightning/contrib/shaper/README.md new file mode 100644 index 000000000..8a3e45e05 --- /dev/null +++ b/contrib/agentlightning/contrib/shaper/README.md @@ -0,0 +1,34 @@ +# SHAPER Runtime Extension + +SHAPER evolves two model-external resources around a frozen Agent Lightning +agent: a textual skill and executable context-construction code. It diagnoses +observable rollout transitions, summarizes episode failures into textual +gradients, and performs a sequential skill-then-harness beam search. + +The runtime lives in `contrib` because Agent Lightning's core resource union +does not yet have a code-harness resource. SHAPER transports both artifacts as +`PromptTemplate` values. Integrations read harness `template` text as source; +they must never call `PromptTemplate.format()` on that source. + +Install Agent Lightning from the repository root and run the SHAPER recipes +from the same checkout: + +```bash +python -m pip install -e . +``` + +Public API: + +- `SHAPER`: the two-stage optimization algorithm. +- `SHAPERTraceAdapter`: extracts structured round and episode records. +- `RoundRecord` and `EpisodeMetadata`: the agent-to-algorithm trace contract. +- `PythonHarnessValidator`: static and isolated-process harness validation. +- `emit_round_record` and `emit_episode_metadata`: rollout instrumentation. + +Generated harnesses are never executed in the Trainer or simulator process. +Validation and every runtime call use the same restricted isolated interpreter +with finite CPU, memory, output, and wall-time limits. This is fault containment, +not an OS sandbox; use a container or VM for code from an untrusted author. + +See [`contrib/recipes/shaper/README.md`](../../../recipes/shaper/README.md) +for VLABench/ESI-Bench environment setup, training, and evaluation commands. diff --git a/contrib/agentlightning/contrib/shaper/__init__.py b/contrib/agentlightning/contrib/shaper/__init__.py new file mode 100644 index 000000000..7b273db3e --- /dev/null +++ b/contrib/agentlightning/contrib/shaper/__init__.py @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""SHAPER: two-stage skill and context-harness evolution for frozen agents.""" + +from .algorithm import ( + DEFAULT_HARNESS_CONTRACT, + SHAPER, + IncomparableCandidateError, + RolloutInfrastructureError, + SkillValidator, + validate_nonempty_skill, +) +from .prompting import parse_json_object +from .roles import ( + ArtifactProposal, + OptimizationStage, + OptimizerRequestContext, + RoleCompleter, + RoleRequest, + SHAPERRoleProtocol, +) +from .sandbox import ( + HarnessOutputValidator, + HarnessRuntimeError, + HarnessValidationResult, + PythonHarnessRuntime, + PythonHarnessValidator, +) +from .trace import SHAPERTraceAdapter, emit_episode_metadata, emit_round_record +from .types import ( + ArtifactCandidate, + ArtifactStage, + CandidateEvaluation, + EpisodeMetadata, + EpisodeSummary, + EpisodeTrace, + OptimizationEvent, + RoundCritique, + RoundRecord, +) + +__all__ = [ + "SHAPER", + "DEFAULT_HARNESS_CONTRACT", + "IncomparableCandidateError", + "RolloutInfrastructureError", + "SkillValidator", + "validate_nonempty_skill", + "ArtifactCandidate", + "ArtifactStage", + "CandidateEvaluation", + "EpisodeMetadata", + "EpisodeSummary", + "EpisodeTrace", + "HarnessValidationResult", + "HarnessOutputValidator", + "HarnessRuntimeError", + "OptimizationEvent", + "PythonHarnessValidator", + "PythonHarnessRuntime", + "RoundCritique", + "RoundRecord", + "SHAPERTraceAdapter", + "ArtifactProposal", + "OptimizationStage", + "OptimizerRequestContext", + "RoleCompleter", + "RoleRequest", + "SHAPERRoleProtocol", + "parse_json_object", + "emit_episode_metadata", + "emit_round_record", +] diff --git a/contrib/agentlightning/contrib/shaper/algorithm.py b/contrib/agentlightning/contrib/shaper/algorithm.py new file mode 100644 index 000000000..9e372ce85 --- /dev/null +++ b/contrib/agentlightning/contrib/shaper/algorithm.py @@ -0,0 +1,969 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Agent Lightning implementation of two-stage SHAPER artifact evolution.""" + +from __future__ import annotations + +import asyncio +import json +import logging +import random +import time +from collections import Counter +from typing import Any, Callable, Dict, Generic, Iterable, Literal, Optional, Sequence, TypeVar, cast + +from openai import AsyncOpenAI + +from agentlightning.adapter import TraceAdapter +from agentlightning.algorithm.base import Algorithm +from agentlightning.store.base import LightningStore +from agentlightning.types import Dataset, NamedResources, PromptTemplate, Rollout, RolloutMode, TaskInput + +from .prompting import load_prompt, parse_json_object +from .roles import OptimizationStage, OptimizerRequestContext, RoleRequest, SHAPERRoleProtocol +from .sandbox import PythonHarnessValidator +from .trace import SHAPERTraceAdapter +from .types import ( + ArtifactCandidate, + CandidateEvaluation, + EpisodeSummary, + EpisodeTrace, + OptimizationEvent, + RoundCritique, + RoundRecord, +) + +logger = logging.getLogger(__name__) + +T_task = TypeVar("T_task") +SkillValidator = Callable[[str], Sequence[str]] + + +class IncomparableCandidateError(RuntimeError): + """A candidate cannot be ranked on the fixed validation denominator.""" + + +class RolloutInfrastructureError(RuntimeError): + """Rollout orchestration failed independently of an artifact's behavior.""" + + +def validate_nonempty_skill(source: str) -> Sequence[str]: + """Apply the benchmark-agnostic minimum contract for planner skills.""" + + return [] if source.strip() else ["Skill must not be empty."] + + +def _diagnostic_context(value: Any) -> tuple[Any, list[dict[str, Any]]]: + """Separate multimodal context images from a compact JSON description.""" + + images: list[dict[str, Any]] = [] + + def visit(item: Any) -> Any: + if isinstance(item, list): + return [visit(child) for child in cast(list[Any], item)] + if not isinstance(item, dict): + return item + mapping = cast(dict[str, Any], item) + if mapping.get("type") == "image_url" and isinstance(mapping.get("image_url"), dict): + image_value = cast(dict[str, Any], mapping["image_url"]) + raw_url = image_value.get("url") + if isinstance(raw_url, str): + images.append(mapping) + return { + "type": "image_url", + "image_url": {"url": ""}, + } + return {str(key): visit(child) for key, child in mapping.items()} + + return visit(value), images + + +DEFAULT_HARNESS_CONTRACT = """The module must define: + def build_context(history): ... + +history is a bounded list of observable round dictionaries. A typical record +contains task_instruction, round_index, planner_response, command, +observation_before, observation_after, action_result, execution_steps, and +runtime_errors. Benchmark integrations may add documented observable fields. +The function must return a deterministic JSON-serializable context payload. +It must not read hidden simulator state, ground-truth answers, or external +state, and its work and output size must remain bounded.""" + + +class _BatchSampler(Generic[T_task]): + """Deterministic epoch sampler used for rollout-gradient minibatches.""" + + def __init__(self, dataset: Sequence[T_task], batch_size: int, seed: int) -> None: + if not dataset: + raise ValueError("Training dataset must not be empty.") + if batch_size < 1: + raise ValueError("gradient_batch_size must be at least 1.") + self._dataset = dataset + self._batch_size = min(batch_size, len(dataset)) + self._random = random.Random(seed) + self._indices: list[int] = [] + self._cursor = 0 + + def next(self) -> list[T_task]: + """Return the next minibatch, reshuffling only at epoch boundaries.""" + + if self._cursor + self._batch_size > len(self._indices): + self._indices = list(range(len(self._dataset))) + self._random.shuffle(self._indices) + self._cursor = 0 + selected = self._indices[self._cursor : self._cursor + self._batch_size] + self._cursor += self._batch_size + return [self._dataset[index] for index in selected] + + +class SHAPER(Algorithm, Generic[T_task]): + """Evolve a textual skill and context-code harness around frozen agents. + + SHAPER follows a fixed two-stage schedule. The first stage updates only the + skill while holding the seed harness fixed. The second stage freezes the + selected skill and updates only the harness. Every proposal is selected by + reward on one fixed validation set, and incumbent candidates remain eligible + in each top-K update. + + Agent Lightning currently has a closed resource union. Both artifacts are + therefore transported as [`PromptTemplate`][agentlightning.PromptTemplate] + resources; harness consumers read ``resource.template`` as Python source and + never call ``format`` on it. + """ + + def __init__( + self, + async_openai_client: AsyncOpenAI, + *, + model: str, + skill_resource_name: str = "skill", + harness_resource_name: str = "harness", + gradient_batch_size: int = 4, + validation_size: Optional[int] = None, + beam_width: int = 3, + branch_factor: int = 2, + skill_rounds: int = 2, + harness_rounds: int = 2, + rollout_batch_timeout: float = 3600.0, + optimizer_temperature: float = 0.7, + role_max_completion_tokens: int = 4096, + role_extra_body: Optional[Dict[str, Any]] = None, + api_retries: int = 3, + artifact_repair_attempts: int = 1, + random_seed: int = 0, + skill_validator: Optional[SkillValidator] = None, + harness_validator: Optional[PythonHarnessValidator] = None, + harness_contract: str = DEFAULT_HARNESS_CONTRACT, + judger_prompt: Optional[str] = None, + summarizer_prompt: Optional[str] = None, + skill_optimizer_prompt: Optional[str] = None, + harness_optimizer_prompt: Optional[str] = None, + role_protocol: Optional[SHAPERRoleProtocol] = None, + ) -> None: + """Initialize SHAPER. + + Args: + async_openai_client: Client used by judger, summarizer, and artifact optimizer roles. + model: One frozen model identifier shared by all evolution roles. + skill_resource_name: Named-resource key containing the textual skill. + harness_resource_name: Named-resource key containing harness Python source. + gradient_batch_size: Rollouts summarized into each textual gradient. + validation_size: Optional fixed subset size; ``None`` uses all validation tasks. + beam_width: Number of incumbents retained after each validation round. + branch_factor: Number of proposals sampled from each beam parent. + skill_rounds: Number of skill-only beam rounds. + harness_rounds: Number of harness-only beam rounds. + rollout_batch_timeout: Wall-clock allowance for one concurrent + rollout wave. The total batch allowance scales with validation + size and Trainer runner count. + optimizer_temperature: Sampling temperature for artifact proposals. + role_max_completion_tokens: Output-token limit for evolution-role calls. + role_extra_body: Provider-specific request fields shared by all + SHAPER role-model calls. + api_retries: Number of attempts for a failed role-model request. + artifact_repair_attempts: Extra optimizer calls after skill or harness + validation failure. + random_seed: Seed for train minibatches and fixed validation subsampling. + skill_validator: Benchmark-owned validator for generated planner skills. + harness_validator: Validator for generated harness source. + harness_contract: Observable input schema and output contract supplied + to the harness optimizer. + judger_prompt: Optional replacement for the bundled round-judger prompt. + summarizer_prompt: Optional replacement for the bundled episode prompt. + skill_optimizer_prompt: Optional replacement skill-optimizer prompt. + harness_optimizer_prompt: Optional replacement harness-optimizer prompt. + role_protocol: Optional benchmark-specific role formatting and + parsing protocol. When supplied, it owns all diagnostic and + optimizer requests while SHAPER retains artifact validation, + candidate evaluation, and selection. + """ + + if skill_resource_name == harness_resource_name: + raise ValueError("Skill and harness resource names must differ.") + for name, value in { + "gradient_batch_size": gradient_batch_size, + "beam_width": beam_width, + "branch_factor": branch_factor, + "role_max_completion_tokens": role_max_completion_tokens, + "api_retries": api_retries, + }.items(): + if value < 1: + raise ValueError(f"{name} must be at least 1.") + if skill_rounds < 0 or harness_rounds < 0 or skill_rounds + harness_rounds < 1: + raise ValueError("At least one non-negative skill or harness round is required.") + if validation_size is not None and validation_size < 1: + raise ValueError("validation_size must be at least 1 when provided.") + if rollout_batch_timeout <= 0: + raise ValueError("rollout_batch_timeout must be positive.") + if artifact_repair_attempts < 0: + raise ValueError("artifact_repair_attempts must not be negative.") + if not harness_contract.strip(): + raise ValueError("harness_contract must not be empty.") + + self.async_openai_client = async_openai_client + self.model = model + self.skill_resource_name = skill_resource_name + self.harness_resource_name = harness_resource_name + self.gradient_batch_size = gradient_batch_size + self.validation_size = validation_size + self.beam_width = beam_width + self.branch_factor = branch_factor + self.skill_rounds = skill_rounds + self.harness_rounds = harness_rounds + self.rollout_batch_timeout = rollout_batch_timeout + self.optimizer_temperature = optimizer_temperature + self.role_max_completion_tokens = role_max_completion_tokens + self.role_extra_body = dict(role_extra_body or {}) + self.api_retries = api_retries + self.artifact_repair_attempts = artifact_repair_attempts + self.random_seed = random_seed + self.skill_validator = skill_validator or validate_nonempty_skill + self.harness_validator = harness_validator or PythonHarnessValidator() + self.harness_contract = harness_contract.strip() + + self._version_counter = 0 + self._seed_resources: Optional[NamedResources] = None + self._seed_harness: Optional[PromptTemplate] = None + self._best_candidate: Optional[ArtifactCandidate] = None + self._validation_cache: dict[str, CandidateEvaluation] = {} + self._optimization_history: list[OptimizationEvent] = [] + + self._judger_prompt = judger_prompt or load_prompt("round_judger.txt") + self._summarizer_prompt = summarizer_prompt or load_prompt("episode_summarizer.txt") + self._skill_optimizer_prompt = skill_optimizer_prompt or load_prompt("skill_optimizer.txt") + self._harness_optimizer_prompt = harness_optimizer_prompt or load_prompt("harness_optimizer.txt") + self._role_protocol = role_protocol + + def get_best_candidate(self) -> ArtifactCandidate: + """Return the best validation candidate encountered across both stages.""" + + if self._best_candidate is None: + raise ValueError("SHAPER has not completed an initial validation.") + return self._best_candidate.model_copy(deep=True) + + def get_best_resources(self) -> NamedResources: + """Return frozen initial resources with the best artifacts installed.""" + + best = self.get_best_candidate() + return self._resources_for_candidate(best) + + def get_optimization_history(self) -> list[OptimizationEvent]: + """Return a detached copy of proposal, validation, and rejection history.""" + + return [event.model_copy(deep=True) for event in self._optimization_history] + + def _get_trace_adapter(self) -> SHAPERTraceAdapter: + adapter: TraceAdapter[Any] = self.get_adapter() + if not isinstance(adapter, SHAPERTraceAdapter): + raise ValueError("SHAPER requires SHAPERTraceAdapter as the Trainer adapter.") + return adapter + + def _initial_artifacts(self) -> tuple[PromptTemplate, PromptTemplate]: + resources = self.get_initial_resources() + if resources is None: + raise ValueError("SHAPER requires initial_resources with skill and harness PromptTemplates.") + skill = resources.get(self.skill_resource_name) + harness = resources.get(self.harness_resource_name) + if not isinstance(skill, PromptTemplate): + raise ValueError(f"Resource {self.skill_resource_name!r} must be a PromptTemplate.") + if not isinstance(harness, PromptTemplate): + raise ValueError(f"Resource {self.harness_resource_name!r} must be a PromptTemplate.") + skill_errors = list(self.skill_validator(skill.template)) + if skill_errors: + raise ValueError("Seed skill failed validation: " + "; ".join(skill_errors)) + harness_validation = self.harness_validator.validate(harness.template) + if not harness_validation.valid: + raise ValueError("Seed harness failed validation: " + "; ".join(harness_validation.errors)) + self._seed_resources = dict(resources) + self._seed_harness = harness + return skill, harness + + def _new_candidate( + self, + *, + skill: PromptTemplate, + harness: PromptTemplate, + stage: Literal["seed", "skill", "harness"], + parent_version: Optional[str] = None, + rationale: str = "", + ) -> ArtifactCandidate: + version = f"shaper-v{self._version_counter:04d}" + self._version_counter += 1 + return ArtifactCandidate( + version=version, + skill=skill, + harness=harness, + stage=stage, + parent_version=parent_version, + rationale=rationale, + ) + + def _resources_for_candidate(self, candidate: ArtifactCandidate) -> NamedResources: + if self._seed_resources is None: + raise ValueError("Initial resources have not been loaded.") + resources = dict(self._seed_resources) + resources[self.skill_resource_name] = candidate.skill + resources[self.harness_resource_name] = candidate.harness + return resources + + @staticmethod + def _materialize_dataset(dataset: Optional[Dataset[T_task]], name: str) -> list[T_task]: + if dataset is None: + raise ValueError(f"{name} dataset is required for SHAPER.") + materialized = [dataset[index] for index in range(len(dataset))] + if not materialized: + raise ValueError(f"{name} dataset must not be empty.") + return materialized + + def _select_fixed_validation(self, dataset: Sequence[T_task]) -> list[T_task]: + if self.validation_size is None or self.validation_size >= len(dataset): + return list(dataset) + rng = random.Random(self.random_seed) + indices = sorted(rng.sample(range(len(dataset)), self.validation_size)) + return [dataset[index] for index in indices] + + def _rollout_batch_allowance(self, rollout_count: int) -> float: + """Scale a per-wave allowance to the number of configured runners.""" + + if rollout_count < 1: + raise ValueError("rollout_count must be at least 1.") + try: + n_runners = self.get_trainer().n_runners + except ValueError: + # Direct unit integrations may attach a store without a Trainer. + n_runners = 1 + n_runners = max(1, n_runners) + waves = (rollout_count + n_runners - 1) // n_runners + return self.rollout_batch_timeout * waves + + async def _evaluate_candidate( + self, + candidate: ArtifactCandidate, + dataset: Sequence[T_task], + mode: Literal["train", "val"], + ) -> CandidateEvaluation: + if mode == "val" and candidate.version in self._validation_cache: + return self._validation_cache[candidate.version] + + store = self.get_store() + resources = self._resources_for_candidate(candidate) + update = await store.update_resources(candidate.version, resources) + queued: list[Rollout] = [] + for task in dataset: + rollout = await store.enqueue_rollout( + input=cast(TaskInput, task), + mode=cast(RolloutMode, mode), + resources_id=update.resources_id, + ) + queued.append(rollout) + + rollout_ids = [rollout.rollout_id for rollout in queued] + batch_allowance = self._rollout_batch_allowance(len(rollout_ids)) + deadline = time.monotonic() + batch_allowance + finished: list[Rollout] = [] + while True: + finished = list(await store.wait_for_rollouts(rollout_ids=rollout_ids, timeout=0.0)) + if len(finished) >= len(rollout_ids): + break + remaining = deadline - time.monotonic() + if remaining <= 0: + break + await asyncio.sleep(min(1.0, max(0.01, remaining))) + + finished_by_id = {rollout.rollout_id: rollout for rollout in finished} + unfinished_ids = [rollout_id for rollout_id in rollout_ids if rollout_id not in finished_by_id] + if unfinished_ids: + sample = ", ".join(unfinished_ids[:3]) + suffix = "" if len(unfinished_ids) <= 3 else ", ..." + raise RolloutInfrastructureError( + f"Candidate {candidate.version} exceeded its {batch_allowance:.1f}s rollout " + f"batch allowance with {len(unfinished_ids)}/{len(rollout_ids)} unfinished " + f"rollout(s): {sample}{suffix}. Training is aborted so unfinished simulator " + "work cannot contaminate a later candidate." + ) + + failed_ids = [rollout.rollout_id for rollout in finished if rollout.status != "succeeded"] + if failed_ids: + sample = ", ".join(failed_ids[:3]) + suffix = "" if len(failed_ids) <= 3 else ", ..." + raise RolloutInfrastructureError( + f"Candidate {candidate.version} had {len(failed_ids)} non-succeeded runner " + f"rollout(s): {sample}{suffix}. Benchmark-owned valid failures must return " + "zero reward and observable episode metadata rather than crash the runner." + ) + + adapter = self._get_trace_adapter() + traces: list[EpisodeTrace] = [] + + for queued_rollout in queued: + completed = finished_by_id.get(queued_rollout.rollout_id) + if completed is None: + raise AssertionError("All rollouts were checked for completion above.") + + spans = await store.query_spans(completed.rollout_id) + trace = adapter.adapt(spans).model_copy( + update={ + "rollout_id": completed.rollout_id, + "task": completed.input, + "status": completed.status, + } + ) + traces.append(trace) + + valid_traces = [trace for trace in traces if not trace.metadata.environment_invalid] + if not valid_traces: + raise IncomparableCandidateError( + f"Candidate {candidate.version} produced no valid {mode} rollouts; " + "simulator-invalid runs cannot define an artifact score." + ) + if mode == "val" and len(valid_traces) != len(traces): + invalid_count = len(traces) - len(valid_traces) + raise IncomparableCandidateError( + f"Candidate {candidate.version} produced {invalid_count} simulator-invalid " + f"validation rollout(s). Every candidate must be scored on the same fixed " + "validation tasks; this candidate is not comparable." + ) + reward_sum = sum(float(trace.final_reward or 0.0) for trace in valid_traces) + evaluation = CandidateEvaluation( + candidate_version=candidate.version, + mode=mode, + requested_rollouts=len(dataset), + finished_rollouts=len(finished_by_id), + valid_rollouts=len(valid_traces), + score=reward_sum / len(valid_traces), + traces=traces, + ) + if mode == "val": + candidate.validation_score = evaluation.score + self._validation_cache[candidate.version] = evaluation + self._record_candidate_score(candidate) + self._consider_historical_best(candidate) + + logger.info( + "[%s] %s score %.4f (%d valid, %d/%d rollouts finished)", + candidate.version, + mode, + evaluation.score, + evaluation.valid_rollouts, + evaluation.finished_rollouts, + evaluation.requested_rollouts, + ) + return evaluation + + def _record_candidate_score(self, candidate: ArtifactCandidate) -> None: + for index in range(len(self._optimization_history) - 1, -1, -1): + event = self._optimization_history[index] + if event.candidate_version == candidate.version: + self._optimization_history[index] = event.model_copy( + update={"validation_score": candidate.validation_score} + ) + break + + def _record_candidate_validation_error(self, candidate: ArtifactCandidate, error: str) -> None: + for index in range(len(self._optimization_history) - 1, -1, -1): + event = self._optimization_history[index] + if event.candidate_version == candidate.version: + self._optimization_history[index] = event.model_copy(update={"validation_error": error}) + break + + def _consider_historical_best(self, candidate: ArtifactCandidate) -> None: + score = candidate.validation_score + if score is None: + return + if self._best_candidate is None or score > cast(float, self._best_candidate.validation_score): + self._best_candidate = candidate.model_copy(deep=True) + logger.info("Historical best is now %s at %.4f", candidate.version, score) + + async def _complete_role(self, role_request: RoleRequest) -> str: + """Execute one role request with shared provider and retry settings.""" + + messages: list[dict[str, Any]] = [ + {"role": "system", "content": role_request.system_prompt}, + {"role": "user", "content": role_request.user_content}, + ] + last_error: Optional[BaseException] = None + for attempt in range(self.api_retries): + try: + request: Dict[str, Any] = { + "model": self.model, + "messages": cast(Any, messages), + "temperature": role_request.temperature, + "max_completion_tokens": self.role_max_completion_tokens, + } + if role_request.response_format == "json_object": + request["response_format"] = {"type": "json_object"} + if self.role_extra_body: + request["extra_body"] = self.role_extra_body + response = cast( + Any, + await self.async_openai_client.chat.completions.create( + **request, + ), + ) + content = response.choices[0].message.content + if not isinstance(content, str) or not content.strip(): + raise ValueError("Role model returned empty content.") + return content.strip() + except (Exception, asyncio.CancelledError) as exc: + if isinstance(exc, asyncio.CancelledError): + raise + last_error = exc + if attempt + 1 < self.api_retries: + await asyncio.sleep(min(2**attempt, 4)) + raise RuntimeError(f"SHAPER role model failed after {self.api_retries} attempts: {last_error}") + + async def _complete_json( + self, + *, + system_prompt: str, + user_content: str | list[dict[str, Any]], + temperature: float, + ) -> Dict[str, Any]: + content = await self._complete_role( + RoleRequest( + system_prompt=system_prompt, + user_content=user_content, + temperature=temperature, + response_format="json_object", + ) + ) + return parse_json_object(content) + + async def _judge_round(self, record: RoundRecord) -> RoundCritique: + context_description, context_images = _diagnostic_context(record.context_payload) + header = { + "round_index": record.round_index, + "task_instruction": record.task_instruction, + "planner_response": record.planner_response, + "command": record.command, + "context_payload": context_description, + "execution_steps": record.execution_steps, + "action_result": record.action_result, + "runtime_errors": record.runtime_errors, + } + content: list[dict[str, Any]] = [ + { + "type": "text", + "text": "ROUND RECORD\n" + json.dumps(header, ensure_ascii=False, default=str), + }, + {"type": "text", "text": "IMAGES ROUTED BY THE CONTEXT HARNESS"}, + *context_images, + {"type": "text", "text": "OBSERVATION BEFORE EXECUTION"}, + *record.observation_before, + {"type": "text", "text": "OBSERVATION AFTER EXECUTION"}, + *record.observation_after, + ] + try: + payload = await self._complete_json( + system_prompt=self._judger_prompt, + user_content=content, + temperature=0.0, + ) + payload["round_index"] = record.round_index + return RoundCritique.model_validate(payload) + except (RuntimeError, ValueError) as exc: + logger.warning("Round judger failed for round %d: %s", record.round_index, exc) + return RoundCritique( + round_index=record.round_index, + progress="unclear", + progress_score=0.0, + observable_change="Judger output unavailable.", + command_assessment="Unknown.", + reasoning_assessment="Unknown.", + context_assessment="Unknown.", + likely_cause=f"Diagnostic failure: {exc}", + suggested_fix="Do not infer an artifact change from this round alone.", + ) + + async def _summarize_episode( + self, + trace: EpisodeTrace, + critiques: Sequence[RoundCritique], + ) -> EpisodeSummary: + total_steps = sum(record.execution_steps for record in trace.rounds) + terminal_context = _diagnostic_context(trace.rounds[-1].context_payload)[0] if trace.rounds else None + task_instruction = trace.rounds[0].task_instruction if trace.rounds else "" + user_payload = { + "rollout_id": trace.rollout_id, + "task_instruction": task_instruction, + "status": trace.status, + "environment_reward": float(trace.final_reward or 0.0), + "environment_invalid": trace.metadata.environment_invalid, + "termination_reason": trace.metadata.termination_reason, + "commands": [record.command for record in trace.rounds], + "execution_steps": total_steps, + "runtime_errors": [ + *trace.metadata.runtime_errors, + *(error for record in trace.rounds for error in record.runtime_errors), + *trace.adapter_errors, + ], + "terminal_context_payload": terminal_context, + "round_critiques": [critique.model_dump(mode="json") for critique in critiques], + } + try: + payload = await self._complete_json( + system_prompt=self._summarizer_prompt, + user_content=json.dumps(user_payload, ensure_ascii=False, default=str), + temperature=0.0, + ) + payload["rollout_id"] = trace.rollout_id + payload["reward"] = float(trace.final_reward or 0.0) + payload["environment_invalid"] = trace.metadata.environment_invalid + return EpisodeSummary.model_validate(payload) + except (RuntimeError, ValueError) as exc: + logger.warning("Episode summarizer failed for %s: %s", trace.rollout_id, exc) + return EpisodeSummary( + rollout_id=trace.rollout_id, + reward=float(trace.final_reward or 0.0), + environment_invalid=trace.metadata.environment_invalid, + instruction_fidelity="Unavailable.", + progress_and_outcome=f"Reward={float(trace.final_reward or 0.0):.3f}.", + repetition_or_recovery="Unavailable.", + decomposition_quality="Unavailable.", + context_effectiveness="Unavailable.", + root_cause=f"Diagnostic failure: {exc}", + actionable_change="Do not infer an artifact change from this episode alone.", + ) + + async def _diagnose_episode(self, trace: EpisodeTrace) -> EpisodeSummary: + critiques = await asyncio.gather(*(self._judge_round(record) for record in trace.rounds)) + return await self._summarize_episode(trace, critiques) + + async def _build_textual_gradient(self, evaluation: CandidateEvaluation) -> Any: + if self._role_protocol is not None: + return await self._role_protocol.build_textual_gradient( + evaluation, + self._complete_role, + ) + + valid_traces = [trace for trace in evaluation.traces if not trace.metadata.environment_invalid] + summaries = await asyncio.gather(*(self._diagnose_episode(trace) for trace in valid_traces)) + rewards = [float(trace.final_reward or 0.0) for trace in valid_traces] + commands = [record.command for trace in valid_traces for record in trace.rounds] + runtime_errors = [ + error + for trace in valid_traces + for error in [ + *trace.metadata.runtime_errors, + *trace.adapter_errors, + *(item for record in trace.rounds for item in record.runtime_errors), + ] + ] + statistics = { + "requested_rollouts": evaluation.requested_rollouts, + "finished_rollouts": evaluation.finished_rollouts, + "valid_rollouts": evaluation.valid_rollouts, + "mean_reward": evaluation.score, + "successful_rollouts": sum(reward >= 1.0 for reward in rewards), + "environment_invalid_rollouts": sum(trace.metadata.environment_invalid for trace in evaluation.traces), + "command_frequencies": dict(Counter(commands).most_common(20)), + "runtime_error_frequencies": dict(Counter(runtime_errors).most_common(20)), + } + return json.dumps( + { + "episode_summaries": [summary.model_dump(mode="json") for summary in summaries], + "aggregate_statistics": statistics, + }, + ensure_ascii=False, + indent=2, + ) + + def _history_text(self) -> str: + events = [event.model_dump(mode="json") for event in self._optimization_history[-30:]] + return json.dumps(events, ensure_ascii=False, indent=2) + + async def _propose_artifact( + self, + *, + parent: ArtifactCandidate, + stage: OptimizationStage, + textual_gradient: Any, + round_index: int, + branch_index: int, + ) -> Optional[ArtifactCandidate]: + common = "" + system_prompt = "" + if self._role_protocol is None: + frozen_note = ( + "The seed harness below is fixed in the skill stage." + if stage == "skill" + else "The selected skill below is fixed in the harness stage." + ) + common = ( + f"ROUND: {round_index}\nBRANCH: {branch_index}\nSTAGE: {stage}\n\n" + f"{frozen_note}\n\nCURRENT SKILL\n=============\n{parent.skill.template}\n\n" + f"CURRENT HARNESS\n===============\n{parent.harness.template}\n\n" + f"ROLLOUT-DERIVED TEXTUAL GRADIENT\n================================\n{textual_gradient}\n\n" + f"OPTIMIZATION HISTORY\n====================\n{self._history_text()}" + ) + if stage == "harness": + common += ( + "\n\nHARNESS VALIDATION CONTRACT\n===========================\n" + self.harness_contract + "\n\n" + f"Define one synchronous {self.harness_validator.function_name} function. " + f"The default smoke arguments are {list(self.harness_validator.smoke_args)!r}. " + "Unsupported imports, dynamic execution, reflection, and arbitrary file I/O are rejected; " + "runtime CPU, memory, output, and wall-clock limits contain expensive work." + ) + system_prompt = self._skill_optimizer_prompt if stage == "skill" else self._harness_optimizer_prompt + + feedback = "" + for repair_index in range(self.artifact_repair_attempts + 1): + try: + if self._role_protocol is None: + payload = await self._complete_json( + system_prompt=system_prompt, + user_content=common + feedback, + temperature=self.optimizer_temperature, + ) + rationale = payload.get("rationale") + artifact = payload.get("new_artifact") + else: + request = self._role_protocol.build_optimizer_request( + OptimizerRequestContext( + parent=parent, + stage=stage, + textual_gradient=textual_gradient, + round_index=round_index, + branch_index=branch_index, + optimization_history=tuple(self._optimization_history), + harness_contract=self.harness_contract, + harness_function_name=self.harness_validator.function_name, + harness_smoke_args=tuple(self.harness_validator.smoke_args), + validation_feedback=feedback, + ) + ) + raw_response = await self._complete_role(request) + proposal = self._role_protocol.parse_optimizer_response(stage, raw_response) + rationale = proposal.rationale + artifact = proposal.artifact + except (RuntimeError, ValueError) as exc: + if isinstance(exc, ValueError) and repair_index < self.artifact_repair_attempts: + feedback = ( + "\n\nOUTPUT PARSE ERROR\n" + + str(exc) + + "\nReturn one complete replacement artifact in the required output format." + ) + continue + self._optimization_history.append( + OptimizationEvent( + round_index=round_index, + stage=stage, + parent_version=parent.version, + rationale="Proposal generation failed.", + validation_error=str(exc), + ) + ) + logger.warning("Artifact proposal failed for %s: %s", parent.version, exc) + return None + if not isinstance(rationale, str) or not isinstance(artifact, str) or not artifact.strip(): + feedback = ( + "\n\nVALIDATION ERROR\n" + "Return non-empty string fields rationale and new_artifact. " + "new_artifact must be the complete replacement artifact, not a diff or wrapper." + ) + continue + + validation_errors: list[str] = [] + if stage == "skill": + validation_errors = list(self.skill_validator(artifact)) + else: + validation = self.harness_validator.validate(artifact) + if not validation.valid: + validation_errors = list(validation.errors) + if validation_errors: + error_text = "; ".join(validation_errors) + self._optimization_history.append( + OptimizationEvent( + round_index=round_index, + stage=stage, + parent_version=parent.version, + rationale=rationale, + validation_error=error_text, + ) + ) + feedback = ( + f"\n\n{stage.upper()} VALIDATION FAILED\n" + + error_text + + "\nRepair the invalid artifact below. Make the smallest change needed to satisfy " + "the validator; do not redesign it, add wrappers, or switch artifact types.\n\n" + "PREVIOUS INVALID ARTIFACT\n=========================\n" + + artifact + + "\n\nEND PREVIOUS INVALID ARTIFACT\n" + "Return the corrected complete artifact while preserving evidence-backed intent." + ) + if repair_index < self.artifact_repair_attempts: + continue + return None + + if stage == "skill": + skill = PromptTemplate(template=artifact.strip(), engine="f-string") + harness = parent.harness + else: + skill = parent.skill + harness = PromptTemplate(template=artifact.strip(), engine="f-string") + candidate = self._new_candidate( + skill=skill, + harness=harness, + stage=stage, + parent_version=parent.version, + rationale=rationale, + ) + self._optimization_history.append( + OptimizationEvent( + round_index=round_index, + stage=stage, + parent_version=parent.version, + candidate_version=candidate.version, + rationale=rationale, + ) + ) + return candidate + self._optimization_history.append( + OptimizationEvent( + round_index=round_index, + stage=stage, + parent_version=parent.version, + rationale="Malformed proposal rejected.", + validation_error=feedback.strip() or "Optimizer returned an invalid artifact payload.", + ) + ) + return None + + async def _generate_children( + self, + *, + beam: Sequence[ArtifactCandidate], + stage: OptimizationStage, + round_index: int, + train_sampler: _BatchSampler[T_task], + ) -> list[ArtifactCandidate]: + children: list[ArtifactCandidate] = [] + for parent in beam: + evaluation = await self._evaluate_candidate(parent, train_sampler.next(), "train") + textual_gradient = await self._build_textual_gradient(evaluation) + proposed = await asyncio.gather( + *( + self._propose_artifact( + parent=parent, + stage=stage, + textual_gradient=textual_gradient, + round_index=round_index, + branch_index=branch_index, + ) + for branch_index in range(self.branch_factor) + ) + ) + children.extend(candidate for candidate in proposed if candidate is not None) + + seen = {candidate.artifact_key() for candidate in beam} + unique: list[ArtifactCandidate] = [] + for candidate in children: + key = candidate.artifact_key() + if key not in seen: + seen.add(key) + unique.append(candidate) + return unique + + async def _select_beam( + self, + candidates: Iterable[ArtifactCandidate], + fixed_validation: Sequence[T_task], + ) -> list[ArtifactCandidate]: + candidate_list = list(candidates) + comparable: list[ArtifactCandidate] = [] + for candidate in candidate_list: + if candidate.validation_score is None: + try: + await self._evaluate_candidate(candidate, fixed_validation, "val") + except IncomparableCandidateError as exc: + self._record_candidate_validation_error(candidate, str(exc)) + logger.warning("Rejecting incomparable candidate %s: %s", candidate.version, exc) + continue + comparable.append(candidate) + comparable.sort( + key=lambda candidate: cast(float, candidate.validation_score), + reverse=True, + ) + if not comparable: + raise RuntimeError("SHAPER beam became empty because no candidate had a comparable validation score.") + return comparable[: self.beam_width] + + async def run( + self, + train_dataset: Optional[Dataset[T_task]] = None, + val_dataset: Optional[Dataset[T_task]] = None, + ) -> None: + """Run hierarchical diagnosis and two-stage top-K artifact evolution.""" + + skill, harness = self._initial_artifacts() + self._get_trace_adapter() + training = self._materialize_dataset(train_dataset, "Training") + validation = self._materialize_dataset(val_dataset, "Validation") + fixed_validation = self._select_fixed_validation(validation) + train_sampler = _BatchSampler(training, self.gradient_batch_size, self.random_seed) + + seed = self._new_candidate(skill=skill, harness=harness, stage="seed") + await self._evaluate_candidate(seed, fixed_validation, "val") + beam: list[ArtifactCandidate] = [seed] + + round_index = 0 + for _ in range(self.skill_rounds): + children = await self._generate_children( + beam=beam, + stage="skill", + round_index=round_index, + train_sampler=train_sampler, + ) + beam = await self._select_beam([*beam, *children], fixed_validation) + round_index += 1 + + selected_skill = beam[0] + if self._seed_harness is None: + raise RuntimeError("Seed harness was not initialized.") + if selected_skill.harness.template != self._seed_harness.template: + raise RuntimeError("Skill stage modified the frozen seed harness.") + + harness_seed = selected_skill + beam = [harness_seed] + frozen_skill_text = harness_seed.skill.template + for _ in range(self.harness_rounds): + children = await self._generate_children( + beam=beam, + stage="harness", + round_index=round_index, + train_sampler=train_sampler, + ) + if any(candidate.skill.template != frozen_skill_text for candidate in children): + raise RuntimeError("Harness stage modified the frozen selected skill.") + beam = await self._select_beam([*beam, *children], fixed_validation) + round_index += 1 + + best = self.get_best_candidate() + store: LightningStore = self.get_store() + await store.update_resources(best.version, self._resources_for_candidate(best)) + logger.info( + "SHAPER complete: best=%s score=%.4f stage=%s", + best.version, + cast(float, best.validation_score), + best.stage, + ) diff --git a/contrib/agentlightning/contrib/shaper/prompting.py b/contrib/agentlightning/contrib/shaper/prompting.py new file mode 100644 index 000000000..d6c4c0e05 --- /dev/null +++ b/contrib/agentlightning/contrib/shaper/prompting.py @@ -0,0 +1,35 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Prompt loading and strict JSON-response parsing for SHAPER.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Dict, cast + +PROMPT_DIR = Path(__file__).parent / "prompts" + + +def load_prompt(name: str) -> str: + """Load a version-controlled SHAPER role prompt.""" + + return (PROMPT_DIR / name).read_text(encoding="utf-8").strip() + + +def parse_json_object(text: str) -> Dict[str, Any]: + """Parse one JSON object, tolerating a single surrounding Markdown fence.""" + + cleaned = text.strip() + if cleaned.startswith("```"): + first_newline = cleaned.find("\n") + if first_newline >= 0: + cleaned = cleaned[first_newline + 1 :] + if cleaned.endswith("```"): + cleaned = cleaned[:-3] + cleaned = cleaned.strip() + + value: Any = json.loads(cleaned) + if not isinstance(value, dict): + raise ValueError("Expected one JSON object from SHAPER role model.") + return cast(Dict[str, Any], value) diff --git a/contrib/agentlightning/contrib/shaper/prompts/episode_summarizer.txt b/contrib/agentlightning/contrib/shaper/prompts/episode_summarizer.txt new file mode 100644 index 000000000..1b34842f1 --- /dev/null +++ b/contrib/agentlightning/contrib/shaper/prompts/episode_summarizer.txt @@ -0,0 +1,26 @@ +ROLE: SHAPER_EPISODE_SUMMARIZER + +Compress one embodied rollout into an actionable textual-gradient record. You +receive observable non-visual execution metadata, the terminal context payload, +round-level judgments, and the environment reward. Identify cross-round +patterns instead of repeating every event. + +Do not turn an environment-invalid rollout into ordinary optimization advice. +Do not treat planner-written observations as ground truth. Keep skill failures +separate from context-harness failures when the evidence permits. + +Return exactly one JSON object with this schema: +{ + "rollout_id": "...", + "reward": 0.0, + "environment_invalid": false, + "instruction_fidelity": "target/entity fidelity across the episode", + "progress_and_outcome": "observable progress and terminal outcome", + "repetition_or_recovery": "loops, progress checks, and recovery behavior", + "decomposition_quality": "quality of subgoals or interface-level commands", + "context_effectiveness": "what context was useful, missing, stale, or harmful", + "root_cause": "most supported systematic cause", + "actionable_change": "smallest reusable artifact change supported by evidence" +} + +No Markdown or surrounding prose. diff --git a/contrib/agentlightning/contrib/shaper/prompts/harness_optimizer.txt b/contrib/agentlightning/contrib/shaper/prompts/harness_optimizer.txt new file mode 100644 index 000000000..ddf9208ca --- /dev/null +++ b/contrib/agentlightning/contrib/shaper/prompts/harness_optimizer.txt @@ -0,0 +1,33 @@ +ROLE: SHAPER_HARNESS_OPTIMIZER + +You evolve the context-code harness of a frozen embodied agent. The planner, +executor, action interface, output parser, and selected textual skill are fixed. +Only the Python context builder may change in this stage. + +The harness decides which observable trajectory records reach the planner and +how they are organized. Use the rollout-derived textual gradient, actual context +payloads, validation feedback, and optimization history. Potential reusable +changes include compact action history, sparse evidence-bearing visual memory, +clear source labels, progress/stagnation detection, bounded recovery guidance, +and preservation of task-critical entities. Add a mechanism only when the +feedback supports it. + +Safety and generalization constraints: +- Return one complete Python module defining the required context function. +- Use only fields and helpers documented in the supplied harness contract. +- Never access ground-truth answers, simulator poses, depth, segmentation, + object metadata, hidden phase state, credentials, network, subprocesses, or + arbitrary files. +- Never use eval, exec, compile, dynamic imports, or mutable global state. +- Do not hard-code task IDs, scene names, validation examples, answer positions, + or benchmark lookup tables. +- Keep work and payload size bounded and deterministic. +- Do not modify the frozen skill. + +Return exactly one JSON object: +{ + "rationale": "evidence-backed changes, preserved strengths, and regression risks", + "new_artifact": "complete replacement Python context-harness module" +} + +No code fences, Markdown, or surrounding prose. diff --git a/contrib/agentlightning/contrib/shaper/prompts/round_judger.txt b/contrib/agentlightning/contrib/shaper/prompts/round_judger.txt new file mode 100644 index 000000000..899de65e5 --- /dev/null +++ b/contrib/agentlightning/contrib/shaper/prompts/round_judger.txt @@ -0,0 +1,26 @@ +ROLE: SHAPER_ROUND_JUDGER + +You diagnose one interaction round from a frozen embodied agent. Ground every +claim in the supplied task, planner output, command, execution metadata, and +observations immediately before and after execution. + +Do not infer hidden simulator state, ground-truth object metadata, or facts that +are not visible in the supplied evidence. Treat model-written reasoning and +context text as fallible. Distinguish planner/skill errors, executor mismatch, +context failures, environment failures, and genuinely ambiguous evidence. + +Return exactly one JSON object with this schema: +{ + "round_index": 0, + "progress": "success | partial | failed | unclear", + "progress_score": 0.0, + "observable_change": "concise before/after comparison", + "command_assessment": "whether the interface-level command was useful and executable", + "reasoning_assessment": "whether the planner used the available evidence coherently", + "context_assessment": "which supplied context helped, distracted, or was missing", + "likely_cause": "most supported cause, without claiming hidden state", + "suggested_fix": "one reusable skill or context change" +} + +The progress score is diagnostic only. The environment reward remains the +authoritative episode outcome. No Markdown or surrounding prose. diff --git a/contrib/agentlightning/contrib/shaper/prompts/skill_optimizer.txt b/contrib/agentlightning/contrib/shaper/prompts/skill_optimizer.txt new file mode 100644 index 000000000..fad3c0d32 --- /dev/null +++ b/contrib/agentlightning/contrib/shaper/prompts/skill_optimizer.txt @@ -0,0 +1,27 @@ +ROLE: SHAPER_SKILL_OPTIMIZER + +You evolve the reusable textual skill of a frozen embodied agent. The planner, +executor, action interface, output parser, and seed context harness are fixed. +Only the skill may change in this stage. + +Use the rollout-derived textual gradient and optimization history to diagnose +systematic behavior. A useful skill can improve scene inspection, exact entity +binding, task decomposition, interface-compatible command forms, progress +verification, recovery after partial or stalled execution, and stopping. Keep +working behavior unless the evidence supports changing it. + +Constraints: +- Return a complete replacement skill, not a diff. +- Preserve the documented action and output contract. +- Do not add task IDs, validation answers, scene-specific lookup tables, hidden + simulator fields, or claims of access to unavailable tools. +- Do not modify or restate the context-harness code. +- Prefer concise reusable policies over a catalogue of training examples. + +Return exactly one JSON object: +{ + "rationale": "evidence-backed changes, preserved strengths, and regression risks", + "new_artifact": "complete replacement textual skill" +} + +No Markdown or surrounding prose. diff --git a/contrib/agentlightning/contrib/shaper/roles.py b/contrib/agentlightning/contrib/shaper/roles.py new file mode 100644 index 000000000..a219133cd --- /dev/null +++ b/contrib/agentlightning/contrib/shaper/roles.py @@ -0,0 +1,78 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Extension points for benchmark-faithful SHAPER role protocols.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Awaitable, Callable, Literal, Protocol, Sequence + +from .types import ArtifactCandidate, CandidateEvaluation, OptimizationEvent + +RoleContent = str | list[dict[str, Any]] +RoleResponseFormat = Literal["json_object", "text"] +OptimizationStage = Literal["skill", "harness"] + + +@dataclass(frozen=True) +class RoleRequest: + """One model call made by a benchmark role protocol.""" + + system_prompt: str + user_content: RoleContent + temperature: float = 0.0 + response_format: RoleResponseFormat = "text" + + +@dataclass(frozen=True) +class ArtifactProposal: + """Parsed replacement artifact returned by an optimizer role.""" + + rationale: str + artifact: str + + +@dataclass(frozen=True) +class OptimizerRequestContext: + """State supplied when a benchmark builds one optimizer request.""" + + parent: ArtifactCandidate + stage: OptimizationStage + textual_gradient: Any + round_index: int + branch_index: int + optimization_history: Sequence[OptimizationEvent] + harness_contract: str + harness_function_name: str + harness_smoke_args: Sequence[Any] + validation_feedback: str = "" + + +RoleCompleter = Callable[[RoleRequest], Awaitable[str]] + + +class SHAPERRoleProtocol(Protocol): + """Benchmark-owned formatting and parsing for SHAPER's model roles.""" + + async def build_textual_gradient( + self, + evaluation: CandidateEvaluation, + complete: RoleCompleter, + ) -> Any: + """Diagnose one development batch using benchmark-specific roles.""" + + ... + + def build_optimizer_request(self, context: OptimizerRequestContext) -> RoleRequest: + """Build the skill- or harness-optimizer request for one proposal.""" + + ... + + def parse_optimizer_response( + self, + stage: OptimizationStage, + response: str, + ) -> ArtifactProposal: + """Extract a complete replacement artifact from the role response.""" + + ... diff --git a/contrib/agentlightning/contrib/shaper/sandbox.py b/contrib/agentlightning/contrib/shaper/sandbox.py new file mode 100644 index 000000000..e529b31cc --- /dev/null +++ b/contrib/agentlightning/contrib/shaper/sandbox.py @@ -0,0 +1,516 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Validation and subprocess execution for generated context harnesses.""" + +from __future__ import annotations + +import ast +import json +import subprocess +import sys +import time +from dataclasses import dataclass, field +from typing import Any, Callable, ClassVar, FrozenSet, Sequence, cast + +from pydantic import BaseModel, ConfigDict, Field + + +class HarnessValidationResult(BaseModel): + """Outcome returned before a generated harness is admitted to a rollout.""" + + model_config = ConfigDict(extra="forbid") + + valid: bool + errors: list[str] = Field(default_factory=list) + duration_seconds: float = Field(ge=0.0) + output_preview: str = "" + + +class HarnessRuntimeError(RuntimeError): + """Raised when an admitted harness fails in its restricted runtime.""" + + +HarnessOutputValidator = Callable[[Any], Sequence[str]] + +# Keep this list intentionally small. Optimizer-generated harnesses normally do +# not need imports; these modules are provided only for deterministic, pure +# transformations of the JSON payload supplied by the benchmark adapter. +SUPPORTED_HARNESS_IMPORTS: FrozenSet[str] = frozenset( + {"collections", "functools", "itertools", "json", "math", "re", "textwrap"} +) + + +def _validate_allowed_imports(allowed_imports: FrozenSet[str]) -> None: + unsupported = sorted(set(allowed_imports) - SUPPORTED_HARNESS_IMPORTS) + if unsupported: + raise ValueError( + "Unsupported harness imports: " + + ", ".join(unsupported) + + ". Supported modules: " + + ", ".join(sorted(SUPPORTED_HARNESS_IMPORTS)) + + "." + ) + + +@dataclass(frozen=True) +class PythonHarnessRuntime: + """Call one harness function only through the isolated worker process. + + This class deliberately does not expose the compiled function in the parent + process. Validation and rollout execution therefore use the same builtins, + import policy, resource limits, and wall-clock timeout. + """ + + source: str + function_name: str = "build_context" + timeout_seconds: float = 2.0 + memory_limit_mb: int = 512 + max_output_chars: int = 8_000_000 + allowed_imports: FrozenSet[str] = frozenset() + output_validator: HarnessOutputValidator | None = None + + def __post_init__(self) -> None: + if not self.function_name.isidentifier(): + raise ValueError("function_name must be a valid Python identifier.") + if self.timeout_seconds <= 0: + raise ValueError("timeout_seconds must be positive.") + if self.memory_limit_mb < 1: + raise ValueError("memory_limit_mb must be at least 1.") + if self.max_output_chars < 1: + raise ValueError("max_output_chars must be at least 1.") + _validate_allowed_imports(self.allowed_imports) + + def __call__(self, *args: Any) -> Any: + """Execute the harness with JSON-serializable arguments.""" + + try: + response = _run_isolated( + source=self.source, + function_name=self.function_name, + args=args, + timeout_seconds=self.timeout_seconds, + memory_limit_mb=self.memory_limit_mb, + max_output_chars=self.max_output_chars, + allowed_imports=self.allowed_imports, + include_output=True, + ) + except subprocess.TimeoutExpired as exc: + raise HarnessRuntimeError(f"Harness execution exceeded {self.timeout_seconds:.2f}s timeout.") from exc + except (TypeError, ValueError) as exc: + raise HarnessRuntimeError(f"Harness arguments are not JSON-serializable: {exc}") from exc + if response.get("valid") is not True: + raw_errors: object = response.get("errors", []) + errors = cast(list[object], raw_errors) if isinstance(raw_errors, list) else [raw_errors] + raise HarnessRuntimeError("; ".join(str(item) for item in errors)) + output = response.get("output") + if self.output_validator is not None: + errors = list(self.output_validator(output)) + if errors: + raise HarnessRuntimeError("; ".join(errors)) + return output + + +def _empty_history_probe() -> tuple[list[dict[str, Any]]]: + return ([],) + + +def _single_history_probe() -> tuple[list[dict[str, Any]]]: + return ( + [ + { + "round_index": 0, + "task_instruction": "sandbox smoke", + "planner_response": "inspect", + "command": "inspect target", + "observation_before": [], + "observation_after": [], + "action_result": {"status": "unknown"}, + "execution_steps": 0, + "runtime_errors": [], + } + ], + ) + + +def _multi_history_probe() -> tuple[list[dict[str, Any]]]: + return ( + [ + { + "round_index": 0, + "task_instruction": "different probe", + "planner_response": "move left", + "command": "move left", + "observation_before": [{"type": "text", "text": "target absent"}], + "observation_after": [{"type": "text", "text": "target visible"}], + "action_result": {"handled": True}, + "execution_steps": 7, + "runtime_errors": [], + }, + { + "round_index": 1, + "task_instruction": "different probe", + "planner_response": "answer", + "command": "answer middle", + "observation_before": [], + "observation_after": [], + "action_result": {"handled": False}, + "execution_steps": 1, + "runtime_errors": ["probe-runtime-error"], + }, + ], + ) + + +@dataclass(frozen=True) +class PythonHarnessValidator: + """Apply static checks and execute several paths in the rollout runtime. + + The worker is a reproducibility and fault-containment layer, not a hardened + operating-system security boundary. Deploy optimizer-generated code in a + container or VM when it is not trusted by the machine owner. + """ + + function_name: str = "build_context" + smoke_args: Sequence[Any] = field(default_factory=_single_history_probe) + additional_smoke_args: Sequence[Sequence[Any]] = field( + default_factory=lambda: (_empty_history_probe(), _multi_history_probe()) + ) + timeout_seconds: float = 2.0 + max_source_chars: int = 100_000 + memory_limit_mb: int = 512 + max_output_chars: int = 8_000_000 + allowed_imports: FrozenSet[str] = frozenset() + output_validator: HarnessOutputValidator | None = None + + _forbidden_names: ClassVar[FrozenSet[str]] = frozenset( + { + "__builtins__", + "__import__", + "breakpoint", + "compile", + "delattr", + "dir", + "eval", + "exec", + "exit", + "getattr", + "globals", + "help", + "input", + "locals", + "open", + "quit", + "setattr", + "vars", + } + ) + _forbidden_attributes: ClassVar[FrozenSet[str]] = frozenset( + { + "chmod", + "glob", + "iterdir", + "mkdir", + "open", + "popen", + "read_bytes", + "read_text", + "remove", + "rename", + "replace", + "resolve", + "rglob", + "rmdir", + "symlink_to", + "system", + "touch", + "unlink", + "write_bytes", + "write_text", + } + ) + + def __post_init__(self) -> None: + if not self.function_name.isidentifier(): + raise ValueError("function_name must be a valid Python identifier.") + if self.timeout_seconds <= 0: + raise ValueError("timeout_seconds must be positive.") + if self.max_source_chars < 1: + raise ValueError("max_source_chars must be at least 1.") + if self.memory_limit_mb < 1: + raise ValueError("memory_limit_mb must be at least 1.") + if self.max_output_chars < 1: + raise ValueError("max_output_chars must be at least 1.") + _validate_allowed_imports(self.allowed_imports) + + def validate(self, source: str) -> HarnessValidationResult: + """Reject unsafe source and smoke multiple execution paths.""" + + started = time.monotonic() + errors = self._static_errors(source) + if errors: + return HarnessValidationResult(valid=False, errors=errors, duration_seconds=time.monotonic() - started) + + previews: list[str] = [] + probes = [tuple(self.smoke_args), *(tuple(args) for args in self.additional_smoke_args)] + for index, args in enumerate(probes): + try: + response = _run_isolated( + source=source, + function_name=self.function_name, + args=args, + timeout_seconds=self.timeout_seconds, + memory_limit_mb=self.memory_limit_mb, + max_output_chars=self.max_output_chars, + allowed_imports=self.allowed_imports, + include_output=self.output_validator is not None, + ) + except (TypeError, ValueError) as exc: + return HarnessValidationResult( + valid=False, + errors=[f"Harness smoke arguments are not JSON-serializable: {exc}"], + duration_seconds=time.monotonic() - started, + ) + except subprocess.TimeoutExpired: + return HarnessValidationResult( + valid=False, + errors=[f"Harness probe {index} exceeded {self.timeout_seconds:.2f}s timeout."], + duration_seconds=time.monotonic() - started, + ) + except HarnessRuntimeError as exc: + return HarnessValidationResult( + valid=False, + errors=[f"Probe {index}: {exc}"], + duration_seconds=time.monotonic() - started, + ) + + if response.get("valid") is not True: + raw_errors: object = response.get("errors", []) + probe_errors = cast(list[object], raw_errors) if isinstance(raw_errors, list) else [raw_errors] + return HarnessValidationResult( + valid=False, + errors=[f"Probe {index}: {item}" for item in probe_errors], + duration_seconds=time.monotonic() - started, + ) + if self.output_validator is not None: + output_errors = list(self.output_validator(response.get("output"))) + if output_errors: + return HarnessValidationResult( + valid=False, + errors=[f"Probe {index}: {item}" for item in output_errors], + duration_seconds=time.monotonic() - started, + ) + previews.append(str(response.get("output_preview", ""))) + + return HarnessValidationResult( + valid=True, + duration_seconds=time.monotonic() - started, + output_preview=" | ".join(previews)[:500], + ) + + def runtime(self, source: str) -> PythonHarnessRuntime: + """Validate source and return its only supported execution interface.""" + + result = self.validate(source) + if not result.valid: + raise ValueError("Invalid context harness: " + "; ".join(result.errors)) + return PythonHarnessRuntime( + source=source, + function_name=self.function_name, + timeout_seconds=self.timeout_seconds, + memory_limit_mb=self.memory_limit_mb, + max_output_chars=self.max_output_chars, + allowed_imports=self.allowed_imports, + output_validator=self.output_validator, + ) + + def _static_errors(self, source: str) -> list[str]: + errors: list[str] = [] + if not source.strip(): + return ["Harness source is empty."] + if len(source) > self.max_source_chars: + return [f"Harness source exceeds {self.max_source_chars} characters."] + + try: + tree = ast.parse(source) + except SyntaxError as exc: + return [f"Syntax error at line {exc.lineno}: {exc.msg}"] + + functions = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))] + matching = [node for node in functions if node.name == self.function_name] + if len(matching) != 1: + errors.append(f"Harness must define exactly one top-level {self.function_name} function.") + elif isinstance(matching[0], ast.AsyncFunctionDef): + errors.append(f"{self.function_name} must be synchronous.") + + for node in ast.walk(tree): + line = getattr(node, "lineno", "?") + if isinstance(node, (ast.ClassDef, ast.Global, ast.Nonlocal, ast.While)): + errors.append(f"{type(node).__name__} is not allowed (line {line}).") + elif isinstance(node, (ast.Import, ast.ImportFrom)): + imported = self._imported_modules(node) + denied = sorted(module for module in imported if module not in self.allowed_imports) + if denied: + errors.append(f"Import is not allowed for {', '.join(denied)} (line {line}).") + for alias in node.names: + imported_name = alias.name + bound_name = alias.asname or imported_name.split(".", maxsplit=1)[0] + if any(part.startswith("_") for part in imported_name.split(".")): + errors.append(f"Private import {imported_name!r} is forbidden at line {line}.") + if bound_name.startswith("_") or bound_name in self._forbidden_names: + errors.append(f"Forbidden import binding {bound_name!r} at line {line}.") + elif isinstance(node, ast.Name): + if node.id in self._forbidden_names or node.id.startswith("__"): + errors.append(f"Forbidden name {node.id!r} at line {line}.") + elif isinstance(node, ast.Attribute): + if node.attr.startswith("_"): + errors.append(f"Private attribute access is forbidden at line {line}.") + elif node.attr in self._forbidden_attributes: + errors.append(f"Forbidden attribute {node.attr!r} at line {line}.") + + return list(dict.fromkeys(errors)) + + @staticmethod + def _imported_modules(node: ast.Import | ast.ImportFrom) -> list[str]: + if isinstance(node, ast.Import): + return [alias.name.split(".", maxsplit=1)[0] for alias in node.names] + if node.module is None: + return [] + return [node.module.split(".", maxsplit=1)[0]] + + +def _run_isolated( + *, + source: str, + function_name: str, + args: Sequence[Any], + timeout_seconds: float, + memory_limit_mb: int, + max_output_chars: int, + allowed_imports: FrozenSet[str], + include_output: bool, +) -> dict[str, Any]: + request = json.dumps( + { + "source": source, + "function_name": function_name, + "args": list(args), + "memory_limit_mb": memory_limit_mb, + "cpu_limit_seconds": max(1, int(timeout_seconds) + 1), + "max_output_chars": max_output_chars, + "allowed_imports": sorted(allowed_imports), + "include_output": include_output, + } + ) + completed = subprocess.run( + [sys.executable, "-I", "-c", _ISOLATED_RUNNER], + input=request, + capture_output=True, + check=False, + text=True, + timeout=timeout_seconds, + ) + try: + decoded: object = json.loads(completed.stdout) + except json.JSONDecodeError as exc: + detail = completed.stderr.strip()[-500:] or f"isolated process exited {completed.returncode}" + raise HarnessRuntimeError(f"Harness worker returned no valid JSON: {detail}") from exc + if not isinstance(decoded, dict): + raise HarnessRuntimeError("Harness worker returned an invalid response.") + response = cast(dict[str, Any], decoded) + if completed.returncode != 0 and response.get("valid") is True: + response = {"valid": False, "errors": [f"Harness worker exited {completed.returncode}."]} + return response + + +_ISOLATED_RUNNER = r""" +import json +import sys +from typing import Any + + +def finish(payload, code=0): + sys.stdout.write(json.dumps(payload)) + raise SystemExit(code) + + +try: + request = json.loads(sys.stdin.read()) + memory_limit_mb = int(request.get("memory_limit_mb", 512)) + cpu_limit_seconds = int(request.get("cpu_limit_seconds", 2)) + try: + import resource + memory_bytes = memory_limit_mb * 1024 * 1024 + if sys.platform.startswith("linux"): + resource.setrlimit(resource.RLIMIT_AS, (memory_bytes, memory_bytes)) + resource.setrlimit(resource.RLIMIT_CPU, (cpu_limit_seconds, cpu_limit_seconds)) + except (ImportError, OSError, ValueError): + pass + + allowed_imports = set(request.get("allowed_imports", [])) + real_import = __import__ + + def limited_import(name, globals=None, locals=None, fromlist=(), level=0): + root = name.split(".", 1)[0] + if root not in allowed_imports: + raise ImportError(f"Import of {root!r} is not allowed") + return real_import(name, globals, locals, fromlist, level) + + safe_builtins = { + "Exception": Exception, + "IndexError": IndexError, + "KeyError": KeyError, + "RuntimeError": RuntimeError, + "TypeError": TypeError, + "ValueError": ValueError, + "abs": abs, + "all": all, + "any": any, + "bool": bool, + "dict": dict, + "enumerate": enumerate, + "filter": filter, + "float": float, + "int": int, + "isinstance": isinstance, + "len": len, + "list": list, + "map": map, + "max": max, + "min": min, + "next": next, + "object": object, + "range": range, + "reversed": reversed, + "round": round, + "set": set, + "slice": slice, + "sorted": sorted, + "str": str, + "sum": sum, + "tuple": tuple, + "zip": zip, + } + if allowed_imports: + safe_builtins["__import__"] = limited_import + + namespace = {"__builtins__": safe_builtins, "Any": Any} + exec(compile(request["source"], "", "exec"), namespace, namespace) + function_name = request["function_name"] + function = namespace.get(function_name) + if not callable(function): + finish({"valid": False, "errors": [f"{function_name} is not callable"]}, 1) + + output = function(*request.get("args", [])) + output_json = json.dumps(output) + max_output_chars = int(request.get("max_output_chars", 8000000)) + if len(output_json) > max_output_chars: + finish({"valid": False, "errors": [f"Harness output exceeds {max_output_chars} characters"]}, 1) + response = {"valid": True, "errors": [], "output_preview": output_json[:500]} + if request.get("include_output"): + response["output"] = output + finish(response, 0) +except BaseException as exc: + if isinstance(exc, (KeyboardInterrupt, SystemExit)): + raise + finish({"valid": False, "errors": [f"{type(exc).__name__}: {exc}"]}, 1) +""" diff --git a/contrib/agentlightning/contrib/shaper/trace.py b/contrib/agentlightning/contrib/shaper/trace.py new file mode 100644 index 000000000..f2669fb37 --- /dev/null +++ b/contrib/agentlightning/contrib/shaper/trace.py @@ -0,0 +1,65 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Trace emission and adaptation helpers for SHAPER rollouts.""" + +from __future__ import annotations + +from typing import Sequence, cast + +from agentlightning.adapter import TraceAdapter +from agentlightning.emitter import emit_object, get_object_value +from agentlightning.reward import find_final_reward +from agentlightning.semconv import AGL_OBJECT +from agentlightning.types import Span + +from .types import EpisodeMetadata, EpisodeTrace, RoundRecord + + +def emit_round_record(record: RoundRecord) -> None: + """Emit one observable planner/executor transition into the active trace.""" + + emit_object(record.model_dump(mode="json"), attributes={"shaper.record_type": "round"}) + + +def emit_episode_metadata(metadata: EpisodeMetadata) -> None: + """Emit optional episode validity and termination metadata.""" + + emit_object(metadata.model_dump(mode="json"), attributes={"shaper.record_type": "episode"}) + + +class SHAPERTraceAdapter(TraceAdapter[EpisodeTrace]): + """Extract SHAPER records and final reward from an Agent Lightning trace.""" + + def adapt(self, source: Sequence[Span], /) -> EpisodeTrace: + rounds: list[RoundRecord] = [] + metadata = EpisodeMetadata() + errors: list[str] = [] + + for span in sorted(source, key=lambda item: item.sequence_id): + if span.name != AGL_OBJECT: + continue + try: + payload: object = get_object_value(span) + except (RuntimeError, TypeError, ValueError) as exc: + errors.append(f"object span {span.span_id}: {exc}") + continue + if not isinstance(payload, dict): + continue + + object_payload = cast(dict[str, object], payload) + record_type = object_payload.get("record_type") + try: + if record_type == "shaper_round": + rounds.append(RoundRecord.model_validate(object_payload)) + elif record_type == "shaper_episode": + metadata = EpisodeMetadata.model_validate(object_payload) + except ValueError as exc: + errors.append(f"invalid {record_type!r} record in span {span.span_id}: {exc}") + + rounds.sort(key=lambda item: item.round_index) + return EpisodeTrace( + final_reward=find_final_reward(source), + rounds=rounds, + metadata=metadata, + adapter_errors=errors, + ) diff --git a/contrib/agentlightning/contrib/shaper/types.py b/contrib/agentlightning/contrib/shaper/types.py new file mode 100644 index 000000000..7c0a842e4 --- /dev/null +++ b/contrib/agentlightning/contrib/shaper/types.py @@ -0,0 +1,145 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Typed records shared by the SHAPER algorithm, agents, and trace adapter.""" + +from __future__ import annotations + +from typing import Any, Dict, List, Literal, Optional + +from pydantic import BaseModel, ConfigDict, Field + +from agentlightning.types import PromptTemplate, RolloutStatus + +ArtifactStage = Literal["seed", "skill", "harness"] +"""Artifact currently being evolved by SHAPER.""" + + +class RoundRecord(BaseModel): + """Observable record for one planner/executor interaction round. + + Agents should emit one record after every executed command. Observations use + OpenAI-compatible content parts so a diagnostic model can consume text, + image URLs, or base64-encoded images without a benchmark-specific adapter. + Hidden simulator state and ground-truth annotations must not be included. + """ + + model_config = ConfigDict(extra="forbid") + + record_type: Literal["shaper_round"] = "shaper_round" + round_index: int = Field(ge=0) + task_instruction: str + planner_response: str + command: str + observation_before: List[Dict[str, Any]] = Field(default_factory=lambda: list[Dict[str, Any]]()) + observation_after: List[Dict[str, Any]] = Field(default_factory=lambda: list[Dict[str, Any]]()) + context_payload: Any = None + harness_input: Any = None + execution_steps: int = Field(default=0, ge=0) + action_result: Dict[str, Any] = Field(default_factory=dict) + runtime_errors: List[str] = Field(default_factory=list) + + +class EpisodeMetadata(BaseModel): + """Optional non-visual metadata emitted once near the end of an episode.""" + + model_config = ConfigDict(extra="forbid") + + record_type: Literal["shaper_episode"] = "shaper_episode" + environment_invalid: bool = False + termination_reason: str = "" + runtime_errors: List[str] = Field(default_factory=list) + extra: Dict[str, Any] = Field(default_factory=dict) + + +class EpisodeTrace(BaseModel): + """Structured trajectory extracted from Agent Lightning spans.""" + + model_config = ConfigDict(extra="forbid") + + rollout_id: str = "" + task: Any = None + status: RolloutStatus = "failed" + final_reward: Optional[float] = None + rounds: List[RoundRecord] = Field(default_factory=lambda: list[RoundRecord]()) + metadata: EpisodeMetadata = Field(default_factory=EpisodeMetadata) + adapter_errors: List[str] = Field(default_factory=list) + + +class RoundCritique(BaseModel): + """Judger output grounded in one before/after execution transition.""" + + model_config = ConfigDict(extra="forbid") + + round_index: int = Field(ge=0) + progress: Literal["success", "partial", "failed", "unclear"] + progress_score: float = Field(ge=0.0, le=1.0) + observable_change: str + command_assessment: str + reasoning_assessment: str + context_assessment: str + likely_cause: str + suggested_fix: str + + +class EpisodeSummary(BaseModel): + """Compact episode-level textual gradient input.""" + + model_config = ConfigDict(extra="forbid") + + rollout_id: str + reward: float + environment_invalid: bool + instruction_fidelity: str + progress_and_outcome: str + repetition_or_recovery: str + decomposition_quality: str + context_effectiveness: str + root_cause: str + actionable_change: str + + +class ArtifactCandidate(BaseModel): + """Versioned pair of model-external artifacts optimized by SHAPER.""" + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + + version: str + skill: PromptTemplate + harness: PromptTemplate + stage: ArtifactStage + parent_version: Optional[str] = None + rationale: str = "" + validation_score: Optional[float] = None + + def artifact_key(self) -> tuple[str, str]: + """Return a stable value key used to remove duplicate proposals.""" + + return self.skill.template, self.harness.template + + +class CandidateEvaluation(BaseModel): + """Reward and trace bundle from evaluating one candidate on one batch.""" + + model_config = ConfigDict(extra="forbid") + + candidate_version: str + mode: Literal["train", "val"] + requested_rollouts: int = Field(ge=0) + finished_rollouts: int = Field(ge=0) + valid_rollouts: int = Field(ge=0) + score: float + traces: List[EpisodeTrace] = Field(default_factory=lambda: list[EpisodeTrace]()) + + +class OptimizationEvent(BaseModel): + """Serializable optimization-history entry supplied to later optimizer calls.""" + + model_config = ConfigDict(extra="forbid") + + round_index: int = Field(ge=0) + stage: Literal["skill", "harness"] + parent_version: str + candidate_version: Optional[str] = None + rationale: str + validation_score: Optional[float] = None + validation_error: Optional[str] = None diff --git a/contrib/recipes/shaper/README.md b/contrib/recipes/shaper/README.md new file mode 100644 index 000000000..d68714e20 --- /dev/null +++ b/contrib/recipes/shaper/README.md @@ -0,0 +1,168 @@ +# SHAPER: Skill-Harness Evolution + +SHAPER evolves a planner skill and an executable context harness around a +frozen embodied agent. The implementation includes runnable VLABench and +ESI-Bench adapters. + +Paper: [arXiv:2608.11350](https://arxiv.org/abs/2608.11350) + +## Install + +Install Agent Lightning from the repository root: + +```bash +python -m pip install -e . +``` + +Check out the benchmark revisions used by the adapters: + +```bash +export AGL_ROOT="$PWD" +export BENCH_ROOT="$AGL_ROOT/shaper-benchmarks" +bash contrib/recipes/shaper/scripts/checkout_shaper_benchmarks.sh "$BENCH_ROOT" +``` + +### VLABench + +VLABench uses separate simulator and OpenPI actor environments. + +```bash +conda create -n shaper-vlabench python=3.10 pip -y +conda activate shaper-vlabench +PYTHON="$CONDA_PREFIX/bin/python" \ + bash "$AGL_ROOT/contrib/recipes/shaper/scripts/bootstrap_shaper_environment.sh" \ + vlabench-simulator "$AGL_ROOT" "$BENCH_ROOT/VLABench" \ + "$BENCH_ROOT/OpenPI" --download-assets +``` + +Install the actor and checkpoint, then start its websocket service: + +```bash +bash "$AGL_ROOT/contrib/recipes/shaper/scripts/bootstrap_shaper_environment.sh" \ + vlabench-actor "$AGL_ROOT" "$BENCH_ROOT/OpenPI" \ + /absolute/path/models/pi0-primitive-10task --download-checkpoint + +bash "$AGL_ROOT/contrib/recipes/shaper/scripts/start_shaper_vlabench_actor.sh" \ + "$AGL_ROOT" "$BENCH_ROOT/OpenPI" \ + /absolute/path/models/pi0-primitive-10task 8000 vlabench-base +``` + +### ESI-Bench + +ESI-Bench uses a controller environment and an isolated Python 3.11 +OmniGibson worker environment. Use a supported 20/30/40-series NVIDIA GPU. + +```bash +python3 -m venv /absolute/path/envs/shaper-esi-controller +PYTHON=/absolute/path/envs/shaper-esi-controller/bin/python \ + bash "$AGL_ROOT/contrib/recipes/shaper/scripts/bootstrap_shaper_environment.sh" \ + esi-controller "$AGL_ROOT" + +conda create -n shaper-esi-worker python=3.11 pip -y +conda activate shaper-esi-worker +export OMNIGIBSON_DATA_PATH=/absolute/path/omnigibson-data +export SHAPER_ACCEPT_NVIDIA_EULA=YES +export SHAPER_ACCEPT_BEHAVIOR_DATASET_TOS=YES +PYTHON="$CONDA_PREFIX/bin/python" \ + bash "$AGL_ROOT/contrib/recipes/shaper/scripts/bootstrap_shaper_environment.sh" \ + esi-worker "$AGL_ROOT" "$BENCH_ROOT/ESI-Bench" \ + "$BENCH_ROOT/BEHAVIOR-1K" --install-behavior +``` + +### Planner Service + +Both recipes use an OpenAI-compatible multimodal planner. To run one locally +with vLLM: + +```bash +python -m pip install "vllm>=0.19.0" +bash contrib/recipes/shaper/scripts/start_shaper_planner_vllm.sh +curl http://127.0.0.1:8001/v1/models +``` + +The Qwen3.6-27B defaults follow its official vLLM and thinking-mode settings: +8-way tensor parallelism, a 262,144-token context, the `qwen3` reasoning +parser, `temperature=1.0`, `top_p=0.95`, `top_k=20`, `min_p=0`, +`presence_penalty=0`, and `repetition_penalty=1`. Set `SHAPER_MODEL`, +`SHAPER_VLLM_TP_SIZE`, `SHAPER_VLLM_MAX_MODEL_LEN`, or the planner sampling +variables in the run scripts when needed. Extra arguments are forwarded to +`vllm serve`. To use a hosted API instead, configure the same run scripts +directly: + +```bash +export SHAPER_PLANNER_ENDPOINT="https://provider.example.com/v1" +export SHAPER_MODEL="" +export OPENAI_API_KEY="" +``` + +## Run + +Runtime configuration is collected at the top of two executable scripts: + +- `contrib/recipes/shaper/scripts/run_vlabench.sh` +- `contrib/recipes/shaper/scripts/run_esi_bench.sh` + +Edit their `Configuration` blocks when your paths or endpoints differ. The +defaults assume benchmark checkouts under `./shaper-benchmarks`, a planner at +`http://127.0.0.1:8001/v1`, and the OpenPI actor at `127.0.0.1:8000`. +`run_esi_bench.sh` also contains the worker-Python and OmniGibson data paths. +Shell environment variables override every value in the scripts. API keys are +never stored in them; export `OPENAI_API_KEY` when the planner requires one. + +Check the configured environment before starting a run: + +```bash +bash contrib/recipes/shaper/scripts/run_vlabench.sh check +bash contrib/recipes/shaper/scripts/run_esi_bench.sh check +``` + +Train: + +```bash +bash contrib/recipes/shaper/scripts/run_vlabench.sh train +bash contrib/recipes/shaper/scripts/run_esi_bench.sh train +``` + +Each training run writes `shaper_run.json`, `best_skill.txt`, and +`best_harness.py` under `outputs/shaper/`. + +The main training options are: + +| Option | Default | Meaning | +|---|---:|---| +| `--n-runners` | `1` | Process-isolated simulator workers | +| `--validation-size` | full split | Fixed validation subset size | +| `--gradient-batch-size` | `4` | Rollouts summarized per optimizer update | +| `--beam-width` | `3` | Candidates retained after validation | +| `--branch-factor` | `2` | Proposals generated per parent | +| `--skill-rounds` | `2` | Skill evolution rounds | +| `--harness-rounds` | `2` | Harness evolution rounds | +| `--rollout-batch-timeout` | `3600` | Seconds allowed per concurrent rollout wave | +| `--role-max-completion-tokens` | planner setting | Judger, summarizer, and optimizer output limit | + +Evaluate the resulting artifact pair on the configured validation split: + +```bash +bash contrib/recipes/shaper/scripts/run_vlabench.sh eval +bash contrib/recipes/shaper/scripts/run_esi_bench.sh eval +``` + +Additional CLI options are forwarded to the underlying command. For example: + +```bash +SHAPER_N_RUNNERS=4 bash contrib/recipes/shaper/scripts/run_vlabench.sh train \ + --beam-width 2 --skill-rounds 3 + +bash contrib/recipes/shaper/scripts/run_esi_bench.sh eval \ + --start-index 0 --limit 20 +``` + +To evaluate the bundled ESI-Bench reporting subset: + +```bash +ESI_VALIDATION_SPLIT="$PWD/contrib/recipes/shaper/esi_bench/splits/reported_eval231.txt" \ + bash contrib/recipes/shaper/scripts/run_esi_bench.sh eval +``` + +Evaluation writes aggregate reward and one record per episode to +`evaluation.json`. diff --git a/contrib/recipes/shaper/__init__.py b/contrib/recipes/shaper/__init__.py new file mode 100644 index 000000000..512a3738b --- /dev/null +++ b/contrib/recipes/shaper/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Runnable SHAPER examples and embodied benchmark integrations.""" diff --git a/contrib/recipes/shaper/cli.py b/contrib/recipes/shaper/cli.py new file mode 100644 index 000000000..40f0d30bd --- /dev/null +++ b/contrib/recipes/shaper/cli.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Shared command helpers for benchmark-specific SHAPER entry points.""" + +from __future__ import annotations + +import os +import sys +from typing import Sequence +from urllib.parse import urlsplit + + +def cli_arguments(argv: Sequence[str] | None) -> list[str]: + """Return explicit arguments or the current process arguments.""" + + return list(argv) if argv is not None else sys.argv[1:] + + +def requests_help(arguments: Sequence[str]) -> bool: + """Return whether an entry point should bypass environment preflight.""" + + return any(argument in {"-h", "--help"} for argument in arguments) + + +def endpoint_socket(endpoint: str) -> tuple[str | None, int | None]: + """Return a best-effort host/port pair for a planner endpoint.""" + + parsed = urlsplit(endpoint) + if not parsed.hostname: + return None, None + if parsed.port is not None: + return parsed.hostname, parsed.port + if parsed.scheme == "https": + return parsed.hostname, 443 + if parsed.scheme == "http": + return parsed.hostname, 80 + return None, None + + +def print_preflight_errors(errors: list[str]) -> int: + """Print actionable preflight failures using a stable CLI format.""" + + for error in errors: + print(f"[missing] {error}") + return 2 if errors else 0 + + +def required_environment(name: str) -> str: + """Read one required environment variable for a direct benchmark CLI.""" + + value = os.environ.get(name, "").strip() + if not value: + raise ValueError(f"Set {name} before running this command.") + return value diff --git a/contrib/recipes/shaper/common.py b/contrib/recipes/shaper/common.py new file mode 100644 index 000000000..5fa448cca --- /dev/null +++ b/contrib/recipes/shaper/common.py @@ -0,0 +1,371 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Shared helpers for SHAPER's embodied benchmark integrations.""" + +from __future__ import annotations + +import ast +import base64 +import importlib +import json +import re +import subprocess +from pathlib import Path +from typing import TYPE_CHECKING, Any, Iterable, Mapping, Sequence, cast + +from openai import OpenAI + +if TYPE_CHECKING: + from agentlightning.types import LLM, NamedResources, Rollout + + +def require_prompt(resources: NamedResources, name: str) -> str: + """Return one prompt resource as unformatted artifact text.""" + + from agentlightning.types import PromptTemplate + + resource = resources.get(name) + if not isinstance(resource, PromptTemplate): + raise TypeError(f"Resource {name!r} must be a PromptTemplate.") + return resource.template + + +def require_llm(resources: NamedResources, name: str, rollout: Rollout) -> LLM: + """Return the concrete planner LLM resource for one rollout.""" + + from agentlightning.types import LLM, AttemptedRollout, ProxyLLM + + resource = resources.get(name) + if not isinstance(resource, LLM): + raise TypeError(f"Resource {name!r} must be an LLM.") + if isinstance(resource, ProxyLLM): + if not isinstance(rollout, AttemptedRollout): + raise ValueError("A ProxyLLM requires an AttemptedRollout before planner use.") + return resource.with_attempted_rollout(rollout) + return resource + + +def openai_client(resource: LLM) -> OpenAI: + """Build a synchronous OpenAI-compatible client from an AGL resource.""" + + return OpenAI( + api_key=resource.api_key or "not-required", + base_url=resource.get_base_url(), + timeout=float(resource.sampling_parameters.get("timeout", 300.0)), + max_retries=int(resource.sampling_parameters.get("max_retries", 2)), + ) + + +def image_data_url(image: Any, *, format_name: str = "PNG") -> str: + """Encode an RGB simulator array with VLABench's OpenCV dependency.""" + + cv2 = cast(Any, importlib.import_module("cv2")) + np = cast(Any, importlib.import_module("numpy")) + array = np.asarray(image) + if array.dtype != np.uint8: + array = np.clip(array, 0, 255).astype(np.uint8) + if array.ndim == 3 and array.shape[2] == 3: + array = cv2.cvtColor(array, cv2.COLOR_RGB2BGR) + elif array.ndim == 3 and array.shape[2] == 4: + array = cv2.cvtColor(array, cv2.COLOR_RGBA2BGRA) + + normalized_format = format_name.strip().lower() + extension = ".jpg" if normalized_format in {"jpg", "jpeg"} else ".png" + ok, buffer = cv2.imencode(extension, array) + if not ok: + raise RuntimeError(f"OpenCV could not encode the simulator image as {extension}.") + encoded = base64.b64encode(buffer.tobytes()).decode("ascii") + mime = "jpeg" if extension == ".jpg" else "png" + return f"data:image/{mime};base64,{encoded}" + + +def path_data_url(path: Path) -> str: + """Encode one official RGB path as an OpenAI image content URL.""" + + suffix = path.suffix.lower() + mime = "image/jpeg" if suffix in {".jpg", ".jpeg"} else "image/png" + encoded = base64.b64encode(path.read_bytes()).decode("ascii") + return f"data:{mime};base64,{encoded}" + + +def image_part(url: str) -> dict[str, Any]: + """Construct one OpenAI-compatible image content part.""" + + return {"type": "image_url", "image_url": {"url": url}} + + +def text_part(text: object) -> dict[str, Any]: + """Construct one OpenAI-compatible text content part.""" + + return {"type": "text", "text": str(text)} + + +def normalize_content(value: Any) -> list[dict[str, Any]]: + """Normalize a harness result into OpenAI-compatible content parts.""" + + if isinstance(value, str): + return [text_part(value)] + if not isinstance(value, list): + raise TypeError("Harness output must be a string or a list of content parts.") + output: list[dict[str, Any]] = [] + for item in cast(list[Any], value): + if not isinstance(item, dict): + raise TypeError("Every multimodal harness item must be a dictionary.") + part = cast(dict[str, Any], item) + if part.get("type") == "text" and isinstance(part.get("text"), str): + output.append(part) + elif part.get("type") == "image_url" and isinstance(part.get("image_url"), dict): + image_url = cast(dict[str, Any], part["image_url"]) + if not isinstance(image_url.get("url"), str): + raise TypeError("image_url.url must be a string.") + output.append(part) + else: + raise TypeError("Harness content parts must be text or image_url blocks.") + return output + + +def validate_multimodal_harness_output(value: Any) -> list[str]: + """Validate the exact observable content shape accepted by both recipes.""" + + if isinstance(value, str): + return [] + if not isinstance(value, list): + return ["Harness output must be a string or a list of content parts."] + errors: list[str] = [] + for index, raw_item in enumerate(cast(list[Any], value)): + if not isinstance(raw_item, dict): + errors.append(f"Content part {index} must be a dictionary.") + continue + item = cast(dict[str, Any], raw_item) + part_type = item.get("type") + if part_type == "text": + if set(item) != {"type", "text"} or not isinstance(item.get("text"), str): + errors.append(f"Text part {index} must contain only string fields type and text.") + elif part_type == "image_url": + image_url = item.get("image_url") + if set(item) != {"type", "image_url"} or not isinstance(image_url, dict): + errors.append(f"Image part {index} must contain only type and image_url.") + continue + image_value = cast(dict[str, Any], image_url) + url = image_value.get("url") + if set(image_value) != {"url"} or not isinstance(url, str): + errors.append(f"Image part {index} must contain exactly one string image_url.url.") + elif not url.startswith("data:image/"): + errors.append(f"Image part {index} must reuse an inline data:image URL.") + else: + errors.append(f"Content part {index} has unsupported type {part_type!r}.") + return errors + + +def strip_thinking(text: str) -> str: + """Remove provider-specific hidden-thinking tags from visible planner output.""" + + return re.sub(r".*?", "", text or "", flags=re.DOTALL | re.IGNORECASE).strip() + + +def first_json_object(text: str) -> dict[str, Any] | None: + """Recover the first JSON object from an otherwise noisy completion.""" + + decoder = json.JSONDecoder() + for index, char in enumerate(text or ""): + if char != "{": + continue + try: + value, _ = decoder.raw_decode(text[index:]) + except json.JSONDecodeError: + continue + if isinstance(value, dict): + return cast(dict[str, Any], value) + return None + + +def completion_text(response: Any) -> tuple[str, str | None]: + """Extract visible content and finish reason from a chat completion.""" + + choice = response.choices[0] + content = choice.message.content + return (content if isinstance(content, str) else ""), str(getattr(choice, "finish_reason", "")) + + +def sanitized_action_result(value: Any) -> dict[str, Any]: + """Keep only official, visibly returned action-result fields.""" + + if not isinstance(value, Mapping): + return {} + mapping = cast(Mapping[str, Any], value) + allowed = { + "handled", + "operation", + "action", + "success", + "error", + "reason", + "object", + "target", + "container", + "physical_state", + "attempts", + "current_stack", + } + return {key: mapping[key] for key in allowed if key in mapping} + + +def load_text(directory: Path, name: str) -> str: + """Read a UTF-8 recipe asset.""" + + return (directory / name).read_text(encoding="utf-8") + + +def ensure_jsonable(parts: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]: + """Detach content parts before crossing the harness subprocess boundary.""" + + return cast(list[dict[str, Any]], json.loads(json.dumps(list(parts)))) + + +def git_revision(path: Path) -> str | None: + """Return the containing Git checkout revision without changing files.""" + + try: + completed = subprocess.run( + ["git", "-C", str(path), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + except (FileNotFoundError, subprocess.SubprocessError): + return None + revision = completed.stdout.strip() + return revision or None + + +def git_tracked_changes(path: Path) -> tuple[str, ...] | None: + """Return tracked files changed from HEAD, or ``None`` outside Git.""" + + try: + completed = subprocess.run( + ["git", "-C", str(path), "diff", "--name-only", "HEAD", "--"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + except (FileNotFoundError, subprocess.SubprocessError): + return None + return tuple(line.strip() for line in completed.stdout.splitlines() if line.strip()) + + +def git_head_file(path: Path, relative_path: str) -> str | None: + """Read one UTF-8 file exactly as recorded by the checkout's HEAD.""" + + try: + completed = subprocess.run( + ["git", "-C", str(path), "show", f"HEAD:{relative_path}"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + except (FileNotFoundError, subprocess.SubprocessError, UnicodeDecodeError): + return None + return completed.stdout + + +def git_gitlink_revision(path: Path, relative_path: str) -> str | None: + """Return the commit recorded for one Git submodule without initializing it.""" + + try: + checkout = subprocess.run( + ["git", "-C", str(path), "rev-parse", "--show-toplevel"], + check=True, + capture_output=True, + text=True, + timeout=5, + ).stdout.strip() + completed = subprocess.run( + ["git", "-C", checkout, "ls-tree", "HEAD", "--", relative_path], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + except (FileNotFoundError, subprocess.SubprocessError): + return None + fields = completed.stdout.strip().split() + if len(fields) < 3 or fields[0] != "160000" or fields[1] != "commit": + return None + return fields[2] + + +def check_python_api( + path: Path, + *, + functions: Mapping[str, Iterable[str]] | None = None, + annotated_classes: Mapping[str, Iterable[str]] | None = None, + class_methods: Mapping[str, Mapping[str, Iterable[str]]] | None = None, +) -> list[str]: + """Check a pinned upstream Python API without importing heavy runtimes.""" + + if not path.is_file(): + return [f"Missing upstream Python module: {path}"] + try: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + except (OSError, SyntaxError) as exc: + return [f"Cannot parse upstream Python module {path}: {exc}"] + + top_level_functions = { + node.name: node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + } + top_level_classes = {node.name: node for node in tree.body if isinstance(node, ast.ClassDef)} + errors: list[str] = [] + for name, required_parameters in (functions or {}).items(): + node = top_level_functions.get(name) + if node is None: + errors.append(f"Upstream API {path}:{name} is missing.") + continue + parameters = {argument.arg for argument in [*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs]} + missing = sorted(set(required_parameters) - parameters) + if missing: + errors.append(f"Upstream API {path}:{name} is missing parameters: {', '.join(missing)}") + + for name, required_fields in (annotated_classes or {}).items(): + node = top_level_classes.get(name) + if node is None: + errors.append(f"Upstream API class {path}:{name} is missing.") + continue + fields = { + child.target.id + for child in node.body + if isinstance(child, ast.AnnAssign) and isinstance(child.target, ast.Name) + } + missing = sorted(set(required_fields) - fields) + if missing: + errors.append(f"Upstream API class {path}:{name} is missing fields: {', '.join(missing)}") + + for class_name, required_methods in (class_methods or {}).items(): + class_node = top_level_classes.get(class_name) + if class_node is None: + errors.append(f"Upstream API class {path}:{class_name} is missing.") + continue + methods = { + node.name: node for node in class_node.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + } + for method_name, required_parameters in required_methods.items(): + method = methods.get(method_name) + if method is None: + errors.append(f"Upstream API {path}:{class_name}.{method_name} is missing.") + continue + parameters = { + argument.arg + for argument in [ + *method.args.posonlyargs, + *method.args.args, + *method.args.kwonlyargs, + ] + } + missing = sorted(set(required_parameters) - parameters) + if missing: + errors.append( + f"Upstream API {path}:{class_name}.{method_name} is missing parameters: " + ", ".join(missing) + ) + return errors diff --git a/contrib/recipes/shaper/esi_bench/__init__.py b/contrib/recipes/shaper/esi_bench/__init__.py new file mode 100644 index 000000000..80a8ba6b9 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""SHAPER integration for the official ESI-Bench active-exploration runner.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .agent import ESIBenchAgent, ESIBenchRuntimeConfig + from .dataset import load_datasets + + +def __getattr__(name: str) -> Any: + """Avoid importing Agent Lightning in the isolated simulator worker.""" + + if name in {"ESIBenchAgent", "ESIBenchRuntimeConfig"}: + from .agent import ESIBenchAgent, ESIBenchRuntimeConfig + + return {"ESIBenchAgent": ESIBenchAgent, "ESIBenchRuntimeConfig": ESIBenchRuntimeConfig}[name] + if name == "load_datasets": + from .dataset import load_datasets + + return load_datasets + raise AttributeError(name) + + +__all__ = ["ESIBenchAgent", "ESIBenchRuntimeConfig", "load_datasets"] diff --git a/contrib/recipes/shaper/esi_bench/agent.py b/contrib/recipes/shaper/esi_bench/agent.py new file mode 100644 index 000000000..36e9ebc49 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/agent.py @@ -0,0 +1,351 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Agent Lightning wrapper around a fresh-process official ESI-Bench run.""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +import signal +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, cast + +from agentlightning.litagent import LitAgent +from agentlightning.types import NamedResources, Rollout +from contrib.agentlightning.contrib.shaper import ( + EpisodeMetadata, + RoundRecord, + emit_episode_metadata, + emit_round_record, +) + +from ..common import require_llm, require_prompt +from ..harness_bridge import HarnessBridgeServer +from .contracts import make_harness_validator + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class ESIBenchRuntimeConfig: + """Runtime configuration shared by one ESI-Bench runner process.""" + + esi_bench_root: Path + behavior_root: Path + questions_jsonl: Path + output_root: Path + omnigibson_data_root: Path | None = None + worker_python: Path = Path(sys.executable) + planner_resource_name: str = "planner_llm" + skill_resource_name: str = "skill" + harness_resource_name: str = "harness" + max_steps: int = 30 + min_steps: int = 3 + confidence_threshold: float = 0.85 + max_new_tokens: int = 32_768 + temperature: float = 1.0 + top_p: float = 0.95 + robot: str = "R1" + episode_timeout_seconds: float = 1800.0 + environment_retries: int = 1 + harness_timeout_seconds: float = 3.0 + harness_memory_limit_mb: int = 768 + harness_max_output_chars: int = 24_000_000 + + def __post_init__(self) -> None: + for name in ("max_steps", "min_steps", "max_new_tokens"): + if int(getattr(self, name)) < 1: + raise ValueError(f"{name} must be positive.") + if self.min_steps > self.max_steps: + raise ValueError("min_steps must not exceed max_steps.") + if not 0.0 <= self.confidence_threshold <= 1.0: + raise ValueError("confidence_threshold must be between zero and one.") + if self.episode_timeout_seconds <= 0: + raise ValueError("episode_timeout_seconds must be positive.") + if self.environment_retries < 0: + raise ValueError("environment_retries must be non-negative.") + + +def _run_token(rollout: Rollout, task_id: str) -> str: + value = f"{rollout.rollout_id}:{task_id}".encode("utf-8") + return hashlib.sha256(value).hexdigest()[:20] + + +def _worker_request( + config: ESIBenchRuntimeConfig, + task: Mapping[str, Any], + *, + endpoint: str, + model: str, + api_key: str | None, + sampling_parameters: Mapping[str, Any], + skill: str, + harness_socket: Path, + harness_token: str, + run_dir: Path, +) -> dict[str, Any]: + """Build the private worker request; this object is sent over stdin only.""" + + return { + "esi_bench_root": str(config.esi_bench_root), + "behavior_root": str(config.behavior_root), + "questions_jsonl": str(config.questions_jsonl), + "run_dir": str(run_dir), + "task": dict(task), + "planner": { + "endpoint": endpoint, + "model": model, + "api_key": api_key, + "sampling_parameters": dict(sampling_parameters), + }, + "skill": skill, + "harness_bridge": { + "socket_path": str(harness_socket), + "token": harness_token, + "timeout_seconds": config.harness_timeout_seconds + 5.0, + "max_response_bytes": config.harness_max_output_chars + 2_000_000, + }, + "runtime": { + "max_steps": int(task.get("max_steps", config.max_steps)), + "min_steps": config.min_steps, + "confidence_threshold": config.confidence_threshold, + "max_new_tokens": config.max_new_tokens, + "temperature": config.temperature, + "top_p": config.top_p, + "robot": config.robot, + "harness_timeout_seconds": config.harness_timeout_seconds, + "harness_memory_limit_mb": config.harness_memory_limit_mb, + "harness_max_output_chars": config.harness_max_output_chars, + }, + } + + +def _terminate_process_group(process: subprocess.Popen[str]) -> None: + """Terminate the simulator and any child processes after a hard timeout.""" + + if process.poll() is not None: + return + if os.name == "posix": + try: + os.killpg(process.pid, signal.SIGTERM) + process.wait(timeout=10) + return + except (OSError, subprocess.TimeoutExpired): + try: + os.killpg(process.pid, signal.SIGKILL) + except OSError: + pass + else: + process.kill() + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + pass + + +def _worker_command(config: ESIBenchRuntimeConfig, response_path: Path) -> list[str]: + """Build the command for the isolated simulator interpreter.""" + + return [ + str(config.worker_python), + "-m", + "contrib.recipes.shaper.esi_bench.worker", + "--response-path", + str(response_path), + ] + + +def _worker_environment(config: ESIBenchRuntimeConfig) -> dict[str, str]: + """Build the isolated worker environment from validated runtime paths.""" + + environment = os.environ.copy() + recipe_checkout = Path(__file__).resolve().parents[4] + omnigibson_checkout = config.behavior_root / "OmniGibson" + existing_pythonpath = environment.get("PYTHONPATH", "") + worker_paths = [str(recipe_checkout), str(omnigibson_checkout)] + if existing_pythonpath: + worker_paths.append(existing_pythonpath) + environment["PYTHONPATH"] = os.pathsep.join(worker_paths) + if config.omnigibson_data_root is not None: + data_root = str(config.omnigibson_data_root.expanduser().resolve()) + environment["ESI_OMNIGIBSON_DATA_ROOT"] = data_root + environment["OMNIGIBSON_DATA_PATH"] = data_root + return environment + + +class ESIBenchAgent(LitAgent[dict[str, Any]]): + """Execute SHAPER artifacts through ESI-Bench's official ``run_one``.""" + + def __init__(self, config: ESIBenchRuntimeConfig) -> None: + super().__init__() + self.config = config + + def _run_worker( + self, + task: Mapping[str, Any], + *, + endpoint: str, + model: str, + api_key: str | None, + sampling_parameters: Mapping[str, Any], + skill: str, + harness: str, + attempt_dir: Path, + ) -> tuple[dict[str, Any], int | None, bool, Path]: + """Run one isolated simulator attempt and return its private response.""" + + attempt_dir.mkdir(parents=True, exist_ok=True) + response_path = attempt_dir / "worker_response.json" + log_path = attempt_dir / "simulator.log" + validator = make_harness_validator( + timeout_seconds=self.config.harness_timeout_seconds, + memory_limit_mb=self.config.harness_memory_limit_mb, + max_output_chars=self.config.harness_max_output_chars, + ) + runtime = validator.runtime(harness) + timed_out = False + with HarnessBridgeServer( + runtime, + max_response_bytes=self.config.harness_max_output_chars + 2_000_000, + ) as bridge: + assert bridge.socket_path is not None + request = _worker_request( + self.config, + task, + endpoint=endpoint, + model=model, + api_key=api_key, + sampling_parameters=sampling_parameters, + skill=skill, + harness_socket=bridge.socket_path, + harness_token=bridge.token, + run_dir=attempt_dir, + ) + command = _worker_command(self.config, response_path) + worker_environment = _worker_environment(self.config) + with log_path.open("w", encoding="utf-8") as log_stream: + process = subprocess.Popen( + command, + stdin=subprocess.PIPE, + stdout=log_stream, + stderr=subprocess.STDOUT, + text=True, + start_new_session=os.name == "posix", + cwd=self.config.esi_bench_root, + env=worker_environment, + ) + try: + process.communicate( + json.dumps(request, ensure_ascii=True), + timeout=self.config.episode_timeout_seconds, + ) + except subprocess.TimeoutExpired: + timed_out = True + _terminate_process_group(process) + + payload: dict[str, Any] = {} + if response_path.is_file(): + try: + value: object = json.loads(response_path.read_text(encoding="utf-8")) + if isinstance(value, dict): + payload = cast(dict[str, Any], value) + except (OSError, json.JSONDecodeError) as exc: + logger.warning("Invalid ESI worker response in %s: %s", attempt_dir, exc) + return payload, process.returncode, timed_out, log_path + + def rollout(self, task: dict[str, Any], resources: NamedResources, rollout: Rollout) -> float: + planner = require_llm(resources, self.config.planner_resource_name, rollout) + skill = require_prompt(resources, self.config.skill_resource_name) + harness = require_prompt(resources, self.config.harness_resource_name) + task_id = str(task.get("task_id", "esi/unknown")) + run_dir = self.config.output_root / _run_token(rollout, task_id) + run_dir.mkdir(parents=True, exist_ok=True) + attempted_logs: list[str] = [] + reason = "worker_failure" + message = "ESI-Bench worker did not produce a result." + environment_invalid = False + + for attempt in range(self.config.environment_retries + 1): + payload, return_code, timed_out, log_path = self._run_worker( + task, + endpoint=planner.get_base_url(), + model=planner.model, + api_key=planner.api_key, + sampling_parameters=planner.sampling_parameters, + skill=skill, + harness=harness, + attempt_dir=run_dir / f"attempt_{attempt + 1}", + ) + attempted_logs.append(str(log_path)) + + if payload.get("ok") is True: + rounds = payload.get("rounds", []) + if isinstance(rounds, list): + for value in cast(list[Any], rounds): + emit_round_record(RoundRecord.model_validate(value)) + metadata = EpisodeMetadata.model_validate(payload.get("metadata", {})) + metadata.extra["simulator_attempts"] = attempt + 1 + metadata.extra["worker_logs"] = attempted_logs + emit_episode_metadata(metadata) + return float(payload.get("reward", 0.0)) + + if timed_out: + reason = "worker_hard_timeout" + message = f"ESI-Bench worker exceeded {self.config.episode_timeout_seconds:.1f}s." + environment_invalid = False + elif return_code is not None and (return_code < 0 or return_code in {134, 139}): + reason = "simulator_process_crash" + message = f"ESI-Bench worker exited with code {return_code}." + environment_invalid = True + elif not payload: + raise RuntimeError( + f"ESI-Bench worker produced no valid response for {task_id}; " + f"exit_code={return_code}, see {log_path}." + ) + else: + reason = str(payload.get("termination_reason", "worker_failure")) + message = str(payload.get("error", f"ESI-Bench worker exited with code {return_code}.")) + environment_invalid = bool(payload.get("environment_invalid", False)) + if payload.get("failure_kind") == "infrastructure": + raise RuntimeError(f"ESI-Bench adapter/upstream failure for {task_id}; see {log_path}: {message}") + if ( + return_code not in {None, 0} + and payload.get("failure_kind") + not in { + "planner", + "artifact", + } + and not environment_invalid + ): + raise RuntimeError( + f"ESI-Bench worker exited unexpectedly for {task_id}; " + f"exit_code={return_code}, see {log_path}: {message}" + ) + if environment_invalid and attempt < self.config.environment_retries: + logger.warning( + "Retrying ESI-Bench rollout %s in a fresh process after environment failure: %s", + task_id, + message, + ) + continue + break + + emit_episode_metadata( + EpisodeMetadata( + environment_invalid=environment_invalid, + termination_reason=reason, + runtime_errors=[message], + extra={ + "task_id": task_id, + "simulator_attempts": len(attempted_logs), + "worker_logs": attempted_logs, + }, + ) + ) + logger.warning("ESI-Bench rollout %s failed: %s", task_id, message) + return 0.0 diff --git a/contrib/recipes/shaper/esi_bench/check_env.py b/contrib/recipes/shaper/esi_bench/check_env.py new file mode 100644 index 000000000..168a76de7 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/check_env.py @@ -0,0 +1,365 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Static ESI-Bench/OmniGibson checks that do not initialize a simulator.""" + +from __future__ import annotations + +import argparse +import ast +import importlib.util +import json +import os +import platform +import socket +import subprocess +import sys +from pathlib import Path +from typing import Sequence, cast + +from ..cli import endpoint_socket +from .contracts import ( + BEHAVIOR_ASSET_VERSION, + OMNIGIBSON_ROBOT_ASSET_VERSION, + check_behavior_source, + check_upstream_source, +) +from .dataset import load_datasets + +RECIPE_DIR = Path(__file__).parent + +_RUNTIME_MODULES = { + "cv2": "opencv-python", + "google.genai": "google-genai", + "numpy": "numpy", + "openai": "openai", + "scipy": "scipy", + "torch": "torch", + "yaml": "PyYAML", +} + + +def absolute_executable(path: Path) -> Path: + """Make an interpreter path absolute without resolving its venv symlink.""" + + return Path(os.path.abspath(str(path.expanduser()))) + + +def check_runtime_modules() -> list[str]: + """Check modules imported eagerly by the pinned official pipeline.""" + + errors: list[str] = [] + for module, package in _RUNTIME_MODULES.items(): + try: + spec = importlib.util.find_spec(module) + except (ImportError, ModuleNotFoundError, ValueError): + spec = None + if spec is None: + errors.append(f"Python module {module} is not importable; install {package} in the behavior environment.") + return errors + + +def check_worker_environment(worker_python: Path, behavior_root: Path) -> list[str]: + """Validate imports in the isolated Isaac/OmniGibson worker interpreter.""" + + executable = absolute_executable(worker_python) + if not executable.is_file(): + return [f"ESI worker Python does not exist: {executable}"] + script = """ +import importlib +import importlib.util +import json +import pathlib +import sys + +behavior_root = pathlib.Path(sys.argv[1]).resolve() +modules = { + "cv2": "opencv-python", + "google.genai": "google-genai", + "numpy": "numpy", + "openai": "openai", + "scipy": "scipy", + "torch": "torch", + "yaml": "PyYAML", +} +errors = [] +for module, package in modules.items(): + try: + spec = importlib.util.find_spec(module) + except (ImportError, ModuleNotFoundError, ValueError) as exc: + errors.append(f"Cannot inspect {module}: {exc}") + continue + if spec is None: + errors.append(f"Python module {module} is not importable; install {package} in the worker environment.") + +try: + spec = importlib.util.find_spec("omnigibson") +except (ImportError, ModuleNotFoundError, ValueError) as exc: + spec = None + errors.append(f"Cannot inspect the omnigibson installation: {exc}") +if spec is None or spec.origin is None: + errors.append("Python package omnigibson is not importable in the worker environment.") +else: + origin = pathlib.Path(spec.origin).resolve() + expected = (behavior_root / "OmniGibson").resolve() + try: + origin.relative_to(expected) + except ValueError: + errors.append(f"omnigibson resolves to {origin}, outside pinned source {expected}.") + +try: + importlib.import_module("contrib.recipes.shaper.esi_bench.worker") +except Exception as exc: + errors.append(f"SHAPER ESI worker bridge is not importable: {type(exc).__name__}: {exc}") + +print(json.dumps({ + "python": f"{sys.version_info.major}.{sys.version_info.minor}", + "errors": errors, +})) +""" + try: + environment = os.environ.copy() + repository_root = Path(__file__).resolve().parents[4] + existing_pythonpath = environment.get("PYTHONPATH", "") + worker_paths = [str(repository_root), str(behavior_root / "OmniGibson")] + if existing_pythonpath: + worker_paths.append(existing_pythonpath) + environment["PYTHONPATH"] = os.pathsep.join(worker_paths) + result = subprocess.run( + [str(executable), "-c", script, str(behavior_root)], + check=False, + capture_output=True, + text=True, + timeout=30, + env=environment, + ) + except (OSError, subprocess.SubprocessError) as exc: + return [f"Cannot run ESI worker Python {executable}: {exc}"] + if result.returncode != 0: + detail = (result.stderr or result.stdout).strip() + return [f"ESI worker Python preflight failed at {executable}: {detail}"] + try: + payload_value: object = json.loads(result.stdout) + except json.JSONDecodeError: + return [f"ESI worker Python returned invalid preflight output: {result.stdout!r}"] + if not isinstance(payload_value, dict): + return [f"ESI worker Python returned an invalid preflight payload: {payload_value!r}"] + payload = cast(dict[str, object], payload_value) + raw_errors = payload.get("errors", []) + if not isinstance(raw_errors, list): + return [f"ESI worker Python returned invalid errors: {raw_errors!r}"] + errors = [str(value) for value in cast(list[object], raw_errors)] + python_version = payload.get("python") + if python_version != "3.11": + errors.append(f"ESI worker requires Python 3.11; found {python_version!r} at {executable}.") + return errors + + +def check_omnigibson_assets(data_root: Path) -> list[str]: + """Check the datasets installed by the pinned BEHAVIOR setup.""" + + required = ( + (data_root / "behavior-1k-assets", BEHAVIOR_ASSET_VERSION), + (data_root / "omnigibson-robot-assets", OMNIGIBSON_ROBOT_ASSET_VERSION), + ) + errors: list[str] = [] + for path, expected_version in required: + if not path.is_dir(): + errors.append(f"Missing OmniGibson dataset payload: {path}") + continue + version_path = path / "VERSION" + try: + version = version_path.read_text(encoding="utf-8").strip() + except OSError: + errors.append(f"Missing OmniGibson dataset version marker: {version_path}") + continue + if version != expected_version: + errors.append( + f"OmniGibson dataset {path.name} version {version!r} does not match pinned {expected_version!r}." + ) + return errors + + +def check_map_generation_patch(path: Path) -> list[str]: + """Verify the exact wall-removal setting required by ESI-Bench.""" + + resolved = path.expanduser().resolve() + if not resolved.is_file(): + return [f"Missing OmniGibson map generator: {resolved}"] + try: + tree = ast.parse(resolved.read_text(encoding="utf-8"), filename=str(resolved)) + except (OSError, SyntaxError) as exc: + return [f"Cannot inspect OmniGibson map generator {resolved}: {exc}"] + assignments: list[ast.AST] = [] + for node in tree.body: + if isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == "NEEDED_STRUCTURE_CATEGORIES" for target in node.targets + ): + assignments.append(node.value) + elif ( + isinstance(node, ast.AnnAssign) + and isinstance(node.target, ast.Name) + and node.target.id == "NEEDED_STRUCTURE_CATEGORIES" + and node.value is not None + ): + assignments.append(node.value) + if not assignments: + return [f"{resolved} does not define NEEDED_STRUCTURE_CATEGORIES."] + final = assignments[-1] + if not isinstance(final, ast.Name) or final.id != "FLOOR_CATEGORIES": + return [ + "ESI-Bench requires NEEDED_STRUCTURE_CATEGORIES = FLOOR_CATEGORIES in " + f"{resolved}; the final assignment is {ast.unparse(final)!r}." + ] + return [] + + +def _gpu_name() -> str | None: + try: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + except (FileNotFoundError, subprocess.SubprocessError): + return None + return result.stdout.splitlines()[0].strip() if result.stdout.strip() else None + + +def check_environment( + *, + root: Path, + behavior_root: Path, + omnigibson_data_root: Path, + questions_jsonl: Path, + train_split: Path, + validation_split: Path, + make_maps_path: Path, + worker_python: Path, + planner_endpoint: str | None, + require_planner: bool, +) -> list[str]: + """Return missing or unsupported prerequisites without launching Isaac.""" + + errors: list[str] = [] + errors.extend(check_upstream_source(root)) + errors.extend(check_behavior_source(behavior_root)) + errors.extend(check_worker_environment(worker_python, behavior_root)) + errors.extend(check_omnigibson_assets(omnigibson_data_root)) + errors.extend(check_map_generation_patch(make_maps_path)) + expected_map_path = behavior_root / "asset_pipeline" / "b1k_pipeline" / "usd_conversion" / "make_maps.py" + if make_maps_path.resolve() != expected_map_path.resolve(): + errors.append( + f"ESI_MAKE_MAPS_PATH must point into the pinned BEHAVIOR checkout: expected {expected_map_path}, " + f"got {make_maps_path}." + ) + if platform.system() != "Linux" or platform.machine() not in {"x86_64", "AMD64"}: + errors.append(f"ESI-Bench requires Linux x86_64; found {platform.system()} {platform.machine()}.") + pipeline = root / "src" / "active_explore" / "pipeline.py" + if not pipeline.is_file(): + errors.append(f"Missing official ESI-Bench runner: {pipeline}") + gpu = _gpu_name() + if gpu is None: + errors.append("No NVIDIA GPU was reported by nvidia-smi.") + elif any(marker in gpu.upper() for marker in ("RTX 50", "B100", "B200", "BLACKWELL")): + errors.append( + f"Official ESI-Bench documents poor rendering on 50-series/Blackwell GPUs; found {gpu}. Use a 20/30/40-series GPU." + ) + try: + train, validation = load_datasets( + questions_jsonl, + train_split, + validation_split, + canonical_root=root / "dataset" / "json_clean", + ) + if not train or not validation: + errors.append("ESI-Bench train and validation splits must both be non-empty.") + except (FileNotFoundError, KeyError, TypeError, ValueError) as exc: + errors.append(str(exc)) + if planner_endpoint: + endpoint_host, endpoint_port = endpoint_socket(planner_endpoint) + if endpoint_host is None or endpoint_port is None: + errors.append(f"SHAPER planner endpoint is not a valid HTTP(S) URL: {planner_endpoint!r}") + else: + try: + with socket.create_connection((endpoint_host, endpoint_port), timeout=2.0): + pass + except OSError as exc: + errors.append(f"Planner endpoint {endpoint_host}:{endpoint_port} is unreachable: {exc}") + elif require_planner: + errors.append("Set SHAPER_PLANNER_ENDPOINT to an OpenAI-compatible chat-completions base URL.") + return errors + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--root", type=Path, default=Path(os.environ.get("ESI_BENCH_ROOT", "."))) + parser.add_argument( + "--behavior-root", + type=Path, + default=Path(os.environ.get("ESI_BEHAVIOR_ROOT", "missing-behavior-root")), + ) + parser.add_argument( + "--omnigibson-data-root", + type=Path, + default=Path( + os.environ.get( + "ESI_OMNIGIBSON_DATA_ROOT", + os.environ.get("OMNIGIBSON_DATA_PATH", "missing-omnigibson-data"), + ) + ), + ) + parser.add_argument("--questions-jsonl", type=Path) + parser.add_argument( + "--train-split", + type=Path, + default=Path(os.environ.get("ESI_TRAIN_SPLIT", RECIPE_DIR / "splits" / "recipe_train10.txt")), + ) + parser.add_argument( + "--validation-split", + type=Path, + default=Path(os.environ.get("ESI_VALIDATION_SPLIT", RECIPE_DIR / "splits" / "recipe_validation10.txt")), + ) + parser.add_argument("--planner-endpoint", default=os.environ.get("SHAPER_PLANNER_ENDPOINT")) + parser.add_argument( + "--worker-python", + type=Path, + default=Path(os.environ.get("ESI_WORKER_PYTHON", sys.executable)), + ) + parser.add_argument("--skip-planner-connect", action="store_true") + parser.add_argument( + "--make-maps-path", + type=Path, + default=Path(os.environ.get("ESI_MAKE_MAPS_PATH", "missing-make-maps.py")), + help="Path to BEHAVIOR-1K asset_pipeline/b1k_pipeline/usd_conversion/make_maps.py.", + ) + args = parser.parse_args(argv) + root = args.root.expanduser().resolve() + questions = ( + args.questions_jsonl.expanduser().resolve() + if args.questions_jsonl is not None + else root / "hf_dataset" / "data" / "questions.jsonl" + ) + errors = check_environment( + root=root, + behavior_root=args.behavior_root.expanduser().resolve(), + omnigibson_data_root=args.omnigibson_data_root.expanduser().resolve(), + questions_jsonl=questions, + train_split=args.train_split.expanduser().resolve(), + validation_split=args.validation_split.expanduser().resolve(), + make_maps_path=args.make_maps_path.expanduser().resolve(), + worker_python=absolute_executable(args.worker_python), + planner_endpoint=str(args.planner_endpoint) if args.planner_endpoint else None, + require_planner=not bool(args.skip_planner_connect), + ) + if errors: + for error in errors: + print(f"[missing] {error}") + return 2 + print("ESI-Bench SHAPER prerequisites passed without launching the simulator.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/contrib/recipes/shaper/esi_bench/contracts.py b/contrib/recipes/shaper/esi_bench/contracts.py new file mode 100644 index 000000000..e0d07d64f --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/contracts.py @@ -0,0 +1,517 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""ESI-Bench-specific artifact contracts shared by optimization and rollout.""" + +from __future__ import annotations + +import ast +import importlib.util +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from contrib.agentlightning.contrib.shaper import PythonHarnessValidator + +from ..common import ( + check_python_api, + git_head_file, + git_revision, + git_tracked_changes, + validate_multimodal_harness_output, +) + +UPSTREAM_REPOSITORY = "https://github.com/ESI-Bench/ESI-Bench" +UPSTREAM_COMMIT = "3c1756396f32b1a90c1f72356a7fde45f418e179" +BEHAVIOR_REPOSITORY = "https://github.com/StanfordVL/BEHAVIOR-1K" +BEHAVIOR_COMMIT = "67ad490856dd465d4606663106f81673fc8bf4e8" +BEHAVIOR_MAP_PATH = "asset_pipeline/b1k_pipeline/usd_conversion/make_maps.py" +BEHAVIOR_ASSET_VERSION = "3.9.0rc7" +OMNIGIBSON_ROBOT_ASSET_VERSION = "3.8.2" + +_ORIGINAL_MAP_SETTING = "NEEDED_STRUCTURE_CATEGORIES = FLOOR_CATEGORIES + WALL_CATEGORIES" +_ESI_MAP_SETTING = "NEEDED_STRUCTURE_CATEGORIES = FLOOR_CATEGORIES" +_GENERIC_JSON_INSTRUCTION = "Return exactly one valid JSON object and nothing else." +_PIPELINE_PATH = Path("src/active_explore/pipeline.py") +_INCLINED_PLANE_PATH = Path("src/active_explore/tasks/physical_dynamics/inclined_plane.py") + + +HARNESS_CONTRACT = """Define exactly `def build_context(records)`. +records is a bounded JSON list containing official observable ESI-Bench data. +Past records contain record_kind, step, visible action/answer/confidence/ +reasoning, sanitized action_result and its official JSON text rendering, an +official RGB full_frame, pixel size/quality, an 8x8 grayscale visual signature, +deterministic pixel-only focus_crops, an optional GRID1000 overlay selected +from the visible task prompt, and optional official extra RGBs. The final +current record also contains max_steps, remaining_steps, the official task +prompt, current RGB derivatives, and official question reference RGB +derivatives. For the pinned inclined-plane post-action call, the final current +record instead has call_kind=auxiliary_post_action and an ordered +observable_sequence of official RGB observations and visible text; prior past +records remain available. Return a bounded JSON-serializable string or list of +OpenAI text/image_url parts. The adapter separately supplies the selected skill +and authoritative task prompt; do not duplicate either. Do not access files, +network, simulator state, camera poses, object metadata, AABBs, depth, +segmentation, rewards, ground-truth answers, task IDs, or mutable external +state.""" + + +def _parse_source(path: Path) -> tuple[ast.Module | None, list[str]]: + try: + return ast.parse(path.read_text(encoding="utf-8"), filename=str(path)), [] + except (OSError, SyntaxError) as exc: + return None, [f"Cannot inspect ESI-Bench model-call contract in {path}: {exc}"] + + +def _top_level_function(tree: ast.Module, name: str) -> ast.FunctionDef | ast.AsyncFunctionDef | None: + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == name: + return node + return None + + +def _generate_json_calls(node: ast.AST) -> list[ast.Call]: + return [ + candidate + for candidate in ast.walk(node) + if isinstance(candidate, ast.Call) + and isinstance(candidate.func, ast.Attribute) + and candidate.func.attr == "generate_json" + ] + + +def _keyword(call: ast.Call, name: str) -> ast.AST | None: + return next((keyword.value for keyword in call.keywords if keyword.arg == name), None) + + +def _is_call_to(node: ast.AST | None, name: str) -> bool: + return isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == name + + +def check_model_call_contract(esi_bench_root: Path) -> list[str]: + """Require every official model call to follow one audited observable path.""" + + errors: list[str] = [] + pipeline_path = esi_bench_root / _PIPELINE_PATH + pipeline_tree, parse_errors = _parse_source(pipeline_path) + errors.extend(parse_errors) + if pipeline_tree is None: + return errors + + pipeline_calls = _generate_json_calls(pipeline_tree) + if len(pipeline_calls) != 2: + errors.append( + f"Pinned ESI-Bench pipeline must contain exactly two model generate_json calls; found {len(pipeline_calls)}." + ) + for function_name in ("force_final_choice", "run_one"): + function = _top_level_function(pipeline_tree, function_name) + if function is None: + errors.append(f"Pinned ESI-Bench pipeline is missing {function_name}().") + continue + calls = _generate_json_calls(function) + if len(calls) != 1: + errors.append( + f"ESI-Bench {function_name}() must contain exactly one generate_json call; found {len(calls)}." + ) + continue + call = calls[0] + if not _is_call_to(_keyword(call, "contents"), "collect_contents"): + errors.append( + f"ESI-Bench {function_name}() must route model contents through the replaceable collect_contents harness." + ) + system_instruction = _keyword(call, "system_instruction") + if not ( + isinstance(system_instruction, ast.Constant) + and isinstance(system_instruction.value, str) + and system_instruction.value == _GENERIC_JSON_INSTRUCTION + ): + errors.append( + f"ESI-Bench {function_name}() must keep the audited generic system instruction so the task skill " + "is injected exactly once by collect_contents." + ) + + task_root = esi_bench_root / "src" / "active_explore" / "tasks" + task_calls: list[tuple[Path, ast.Module, list[ast.Call]]] = [] + if not task_root.is_dir(): + errors.append(f"Missing official ESI-Bench task source directory: {task_root}") + return errors + for task_path in sorted(task_root.rglob("*.py")): + task_tree, task_parse_errors = _parse_source(task_path) + errors.extend(task_parse_errors) + if task_tree is None: + continue + calls = _generate_json_calls(task_tree) + if calls: + task_calls.append((task_path.relative_to(esi_bench_root), task_tree, calls)) + + if len(task_calls) != 1 or task_calls[0][0] != _INCLINED_PLANE_PATH or len(task_calls[0][2]) != 1: + locations = [f"{path} ({len(calls)} call(s))" for path, _tree, calls in task_calls] + errors.append( + "ESI-Bench task modules may contain only the audited inclined-plane post-action model call; found: " + + (", ".join(locations) if locations else "none") + + "." + ) + return errors + + inclined_path, inclined_tree, inclined_calls = task_calls[0] + post_action = _top_level_function(inclined_tree, "post_action_query") + if post_action is None or inclined_calls[0] not in _generate_json_calls(post_action): + errors.append(f"{inclined_path} must keep its sole generate_json call inside post_action_query().") + return errors + call = inclined_calls[0] + contents = _keyword(call, "contents") + if not isinstance(contents, ast.Name) or contents.id != "contents": + errors.append("Inclined-plane post_action_query() must send its official observable frame contents.") + if not _is_call_to(_keyword(call, "system_instruction"), "build_system_prompt"): + errors.append("Inclined-plane post_action_query() must use its official task-specific system prompt.") + return errors + + +def check_upstream_source(esi_bench_root: Path) -> list[str]: + """Validate the pinned official runner contract without importing Isaac.""" + + errors: list[str] = [] + revision = git_revision(esi_bench_root) + if revision is None: + errors.append(f"ESI-Bench source is not inside a readable Git checkout: {esi_bench_root}") + elif revision != UPSTREAM_COMMIT: + errors.append(f"ESI-Bench revision {revision} does not match pinned {UPSTREAM_COMMIT}.") + changes = git_tracked_changes(esi_bench_root) + if changes: + errors.append("ESI-Bench checkout has tracked modifications: " + ", ".join(changes) + ".") + errors.extend( + check_python_api( + esi_bench_root / _PIPELINE_PATH, + functions={ + "build_model_client": {"provider", "api_key", "model"}, + "collect_contents": { + "image_path", + "history", + "prompt", + "reference_image_paths", + "reference_image_path", + }, + "force_final_choice": { + "task_module", + "model_client", + "payload", + "camera_info", + "image_path", + "history", + "config", + "task_state", + "reference_image_paths", + }, + "run_one": {"config"}, + }, + annotated_classes={ + "ActiveExploreConfig": { + "task", + "metadata", + "question_index", + "json_root", + "results_root", + "step_image_root", + "provider", + "model", + "api_key", + "max_steps", + "min_steps", + "threshold", + "max_new_tokens", + "temperature", + "top_p", + "robot", + "overwrite", + } + }, + ) + ) + errors.extend(check_model_call_contract(esi_bench_root)) + return errors + + +def check_behavior_source(behavior_root: Path) -> list[str]: + """Validate the contrib deployment pin for BEHAVIOR/OmniGibson.""" + + errors: list[str] = [] + revision = git_revision(behavior_root) + if revision is None: + errors.append(f"BEHAVIOR-1K source is not inside a readable Git checkout: {behavior_root}") + elif revision != BEHAVIOR_COMMIT: + errors.append(f"BEHAVIOR-1K revision {revision} does not match pinned {BEHAVIOR_COMMIT}.") + changes = git_tracked_changes(behavior_root) + unexpected = sorted(set(changes or ()) - {BEHAVIOR_MAP_PATH}) + if unexpected: + errors.append( + "BEHAVIOR-1K checkout has tracked modifications beyond the official map setting: " + + ", ".join(unexpected) + + "." + ) + if changes and BEHAVIOR_MAP_PATH in changes: + baseline = git_head_file(behavior_root, BEHAVIOR_MAP_PATH) + map_path = behavior_root / BEHAVIOR_MAP_PATH + try: + current = map_path.read_text(encoding="utf-8") + except OSError as exc: + errors.append(f"Cannot read patched BEHAVIOR map generator {map_path}: {exc}") + else: + if baseline is None or baseline.count(_ORIGINAL_MAP_SETTING) != 1: + errors.append("Cannot reconstruct the expected ESI map setting from pinned BEHAVIOR HEAD.") + else: + expected = baseline.replace(_ORIGINAL_MAP_SETTING, _ESI_MAP_SETTING, 1) + if current != expected: + errors.append( + "BEHAVIOR map generator differs from pinned HEAD beyond the one ESI-required " + "FLOOR_CATEGORIES setting." + ) + required = ( + behavior_root / "setup.sh", + behavior_root / "OmniGibson" / "omnigibson" / "__init__.py", + behavior_root / BEHAVIOR_MAP_PATH, + ) + errors.extend(f"Missing pinned BEHAVIOR-1K source path: {path}" for path in required if not path.is_file()) + return errors + + +def check_omnigibson_install(behavior_root: Path) -> list[str]: + """Require the active environment to import OmniGibson from the pin.""" + + try: + spec = importlib.util.find_spec("omnigibson") + except (ImportError, ModuleNotFoundError, ValueError) as exc: + return [f"Cannot inspect the omnigibson installation: {exc}"] + if spec is None or spec.origin is None: + return ["Python package omnigibson is not importable."] + origin = Path(spec.origin).resolve() + expected = (behavior_root / "OmniGibson").resolve() + try: + origin.relative_to(expected) + except ValueError: + return [f"omnigibson resolves to {origin}, outside pinned source {expected}."] + return [] + + +def validate_skill(source: str) -> list[str]: + """Preserve the official prompt inside a plain planner-policy artifact.""" + + errors: list[str] = [] + count = source.count("{task_prompt}") + if count != 1: + errors.append(f"Skill must contain {{task_prompt}} exactly once; found {count}.") + if len(source) > 30_000: + errors.append("Skill must remain below 30,000 characters.") + code_prefixes = ("import ", "from ", "def ", "class ", "async def ") + if any(line.lstrip().startswith(code_prefixes) for line in source.splitlines()): + errors.append("Skill must be planner instructions, not Python source.") + lowered = source.lower() + if any(marker in lowered for marker in ("etl_definition", "create_etl_artifact", "artifact[")): + errors.append("Skill must not contain artifact-framework or ETL wrappers.") + if "\x00" in source: + errors.append("Skill must not contain NUL bytes.") + return errors + + +def _image(label: str) -> dict[str, Any]: + return { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{label}"}, + } + + +def _observation(label: str, *, include_grid: bool = False) -> dict[str, Any]: + value: dict[str, Any] = { + "source": "official_rgb", + "pixel_size": {"width": 640, "height": 480}, + "pixel_quality": { + "contrast": 0.6, + "edge_density": 0.15, + "laplacian_sharpness": 0.25, + "score": 1.4375, + }, + "visual_signature": { + "kind": "8x8_grayscale", + "values": [index * 4 for index in range(64)], + }, + "full_frame": _image(f"{label}_FULL"), + "focus_crops": [ + { + "source": "deterministic_pixel_crop", + "region": "horizontal_band_3", + "quality": { + "contrast": 0.4, + "edge_density": 0.2, + "laplacian_sharpness": 0.1, + "score": 1.25, + }, + "image": _image(f"{label}_CENTER"), + }, + { + "source": "deterministic_pixel_crop", + "region": "overlapping_tile_7", + "quality": { + "contrast": 0.5, + "edge_density": 0.1, + "laplacian_sharpness": 0.2, + "score": 1.0, + }, + "image": _image(f"{label}_LOWER"), + }, + ], + } + if include_grid: + value["grid_overlay"] = { + "source": "deterministic_pixel_overlay", + "coordinate_system": "GRID1000: x=0..1000 left-to-right, y=0..1000 top-to-bottom", + "image": _image(f"{label}_GRID1000"), + } + return value + + +def _one_step_probe() -> tuple[list[dict[str, Any]]]: + return ( + [ + { + "record_kind": "past", + "step": 1, + "action": "turn_left", + "answer": "not sure", + "confidence": 0.1, + "reasoning": "Seek another viewpoint.", + "action_result": {"handled": True, "operation": "camera"}, + "action_result_text": '{"handled": true, "operation": "camera"}', + "observation": _observation("PAST_1"), + "extra_observations": [], + }, + { + "record_kind": "current", + "step": 2, + "max_steps": 30, + "remaining_steps": 29, + "task_instruction": "Which candidate matches the reflected object?", + "observation": _observation("CURRENT_2"), + "reference_observations": [ + { + "label": "QUESTION REFERENCE IMAGE 1", + **_observation("REFERENCE_1"), + } + ], + }, + ], + ) + + +def _two_step_probe() -> tuple[list[dict[str, Any]]]: + records = list(_one_step_probe()[0][:-1]) + records.append( + { + "record_kind": "past", + "step": 2, + "action": "move_closer", + "answer": "not sure", + "confidence": 0.4, + "reasoning": "Compare the candidates with retained evidence.", + "action_result": {"handled": True, "operation": "navigation"}, + "action_result_text": '{"handled": true, "operation": "navigation"}', + "observation": _observation("PAST_2"), + "extra_observations": [_observation("EXTRA_2")], + } + ) + records.append( + { + "record_kind": "current", + "step": 3, + "max_steps": 30, + "remaining_steps": 28, + "task_instruction": "Which candidate matches the reflected object?", + "observation": _observation("CURRENT_3"), + "reference_observations": [], + } + ) + return (records,) + + +def _geometry_probe() -> tuple[list[dict[str, Any]]]: + records = list(_one_step_probe()[0][:-1]) + records.append( + { + "record_kind": "current", + "step": 2, + "max_steps": 30, + "remaining_steps": 29, + "task_instruction": "Do the three objects form a straight line?", + "observation": _observation("GEOMETRY_CURRENT", include_grid=True), + "reference_observations": [], + } + ) + return (records,) + + +def _auxiliary_probe() -> tuple[list[dict[str, Any]]]: + records = list(_one_step_probe()[0][:-1]) + records.append( + { + "record_kind": "current", + "call_kind": "auxiliary_post_action", + "step": 2, + "max_steps": 30, + "remaining_steps": 29, + "task_instruction": "Estimate the visible inclined-plane outcome from the frame sequence.", + "observable_sequence": [ + { + "content_kind": "observation", + "sequence_index": 0, + "observation": _observation("AUXILIARY_FRAME_1"), + }, + { + "content_kind": "observation", + "sequence_index": 1, + "observation": _observation("AUXILIARY_FRAME_2"), + }, + { + "content_kind": "text", + "sequence_index": 2, + "text": "Use only the visible frame sequence.", + }, + ], + "reference_observations": [], + } + ) + return (records,) + + +def make_harness_validator( + *, + timeout_seconds: float = 3.0, + memory_limit_mb: int = 768, + max_output_chars: int = 24_000_000, +) -> PythonHarnessValidator: + """Build the validator used both when admitting and executing artifacts.""" + + from contrib.agentlightning.contrib.shaper import PythonHarnessValidator + + return PythonHarnessValidator( + smoke_args=_one_step_probe(), + additional_smoke_args=(([],), _two_step_probe(), _geometry_probe(), _auxiliary_probe()), + timeout_seconds=timeout_seconds, + memory_limit_mb=memory_limit_mb, + max_output_chars=max_output_chars, + output_validator=validate_multimodal_harness_output, + ) + + +__all__ = [ + "BEHAVIOR_ASSET_VERSION", + "BEHAVIOR_COMMIT", + "BEHAVIOR_MAP_PATH", + "BEHAVIOR_REPOSITORY", + "HARNESS_CONTRACT", + "OMNIGIBSON_ROBOT_ASSET_VERSION", + "UPSTREAM_COMMIT", + "UPSTREAM_REPOSITORY", + "check_behavior_source", + "check_model_call_contract", + "check_omnigibson_install", + "check_upstream_source", + "make_harness_validator", + "validate_skill", +] diff --git a/contrib/recipes/shaper/esi_bench/dataset.py b/contrib/recipes/shaper/esi_bench/dataset.py new file mode 100644 index 000000000..0d32a7d48 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/dataset.py @@ -0,0 +1,212 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Explicit ESI-Bench split loading without exposing labels to the planner.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Iterable, Mapping, cast + + +def normalize_question_id(value: object) -> str: + """Normalize HF and runner question identifiers to four digit strings.""" + + text = str(value).strip() + if text.lower().startswith("q_"): + text = text[2:] + if not text.isdigit(): + raise ValueError(f"Invalid ESI-Bench question id: {value!r}") + return text.zfill(4) + + +def read_split(path: Path) -> list[str]: + """Read one ordered, duplicate-free question-id manifest.""" + + if not path.is_file(): + raise FileNotFoundError(f"ESI-Bench split does not exist: {path}") + ids = [normalize_question_id(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()] + if not ids: + raise ValueError(f"ESI-Bench split is empty: {path}") + if len(ids) != len(set(ids)): + raise ValueError(f"ESI-Bench split contains duplicate ids: {path}") + return ids + + +def load_question_rows(path: Path) -> dict[str, dict[str, Any]]: + """Index the official Hugging Face JSONL export by normalized id.""" + + if not path.is_file(): + raise FileNotFoundError(f"Official ESI-Bench questions.jsonl not found: {path}") + rows: dict[str, dict[str, Any]] = {} + with path.open("r", encoding="utf-8") as stream: + for line_number, line in enumerate(stream, start=1): + if not line.strip(): + continue + value: object = json.loads(line) + if not isinstance(value, dict): + raise ValueError(f"Expected a JSON object at {path}:{line_number}") + row = cast(dict[str, Any], value) + question_id = normalize_question_id(row.get("id")) + if question_id in rows: + raise ValueError(f"Duplicate ESI-Bench question id {question_id} in {path}") + rows[question_id] = row + return rows + + +def load_question_row(path: Path, question_id: object) -> dict[str, Any]: + """Load one scorer-side HF row without placing labels in an AGL task.""" + + normalized = normalize_question_id(question_id) + if not path.is_file(): + raise FileNotFoundError(f"Official ESI-Bench questions.jsonl not found: {path}") + with path.open("r", encoding="utf-8") as stream: + for line_number, line in enumerate(stream, start=1): + if not line.strip(): + continue + value: object = json.loads(line) + if not isinstance(value, dict): + raise ValueError(f"Expected a JSON object at {path}:{line_number}") + row = cast(dict[str, Any], value) + if normalize_question_id(row.get("id")) == normalized: + return row + raise KeyError(f"Question {normalized} is absent from the official ESI-Bench JSONL export.") + + +def index_canonical_questions(json_root: Path) -> dict[str, Path]: + """Index the official runner JSON files while ignoring aggregate manifests.""" + + resolved_root = json_root.expanduser().resolve() + if not resolved_root.is_dir(): + raise FileNotFoundError(f"Official ESI-Bench dataset/json_clean directory not found: {resolved_root}") + paths: dict[str, Path] = {} + for path in sorted(resolved_root.rglob("*.json")): + value: object = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict) or "id" not in value: + continue + row = cast(dict[str, Any], value) + question_id = normalize_question_id(row["id"]) + previous = paths.get(question_id) + if previous is not None: + raise ValueError(f"Duplicate canonical ESI-Bench question {question_id}: {previous} and {path}") + if not str(row.get("runner_task", "")).strip(): + raise ValueError(f"Canonical ESI-Bench question {question_id} has no runner_task: {path}") + paths[question_id] = path.resolve() + if not paths: + raise ValueError(f"No canonical ESI-Bench question files found below {resolved_root}") + return paths + + +def resolve_canonical_question(json_root: Path, relative_path: object, question_id: object) -> Path: + """Resolve a task's canonical runner JSON without allowing path traversal.""" + + root = json_root.expanduser().resolve() + normalized = normalize_question_id(question_id) + if not isinstance(relative_path, str) or not relative_path.strip(): + path = index_canonical_questions(root).get(normalized) + if path is None: + raise KeyError(f"Question {normalized} is absent from canonical ESI-Bench JSON files.") + else: + path = (root / relative_path).resolve() + try: + path.relative_to(root) + except ValueError as exc: + raise ValueError(f"Canonical question path escapes dataset/json_clean: {path}") from exc + if not path.is_file(): + raise FileNotFoundError(f"Canonical ESI-Bench question does not exist: {path}") + value: object = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict) or normalize_question_id(cast(dict[str, Any], value).get("id")) != normalized: + raise ValueError(f"Canonical ESI-Bench question id mismatch for {path}; expected {normalized}.") + return path + + +def materialize_split( + rows: Mapping[str, dict[str, Any]], + question_ids: Iterable[str], + *, + max_steps: int, + canonical_questions: Mapping[str, Path] | None = None, + canonical_root: Path | None = None, +) -> list[dict[str, Any]]: + """Build tasks consumed by the official runner. + + The AGL task contains identifiers and runner configuration only. The worker + loads the scorer-owned HF row by ID after it has crossed into the isolated + simulator process. Labels never enter planner or harness inputs; after the + episode finishes, the worker may attach them to post-hoc diagnostic + metadata for the Judger described by the SHAPER protocol. + """ + + tasks: list[dict[str, Any]] = [] + for raw_id in question_ids: + question_id = normalize_question_id(raw_id) + row = rows.get(question_id) + if row is None: + raise KeyError(f"Question {question_id} is absent from the official ESI-Bench JSONL export.") + runner_task = str(row.get("runner_task", "")).strip() + if not runner_task: + raise ValueError(f"Question {question_id} has no runner_task.") + task: dict[str, Any] = { + "task_id": f"esi/{question_id}", + "question_id": question_id, + "runner_task": runner_task, + "max_steps": max_steps, + } + if canonical_questions is not None: + if canonical_root is None: + raise ValueError("canonical_root is required with canonical_questions.") + canonical_path = canonical_questions.get(question_id) + if canonical_path is None: + raise KeyError(f"Question {question_id} is absent from canonical ESI-Bench JSON files.") + canonical_value: object = json.loads(canonical_path.read_text(encoding="utf-8")) + if not isinstance(canonical_value, dict): + raise ValueError(f"Canonical ESI-Bench question must be an object: {canonical_path}") + canonical_row = cast(dict[str, Any], canonical_value) + if str(canonical_row.get("runner_task", "")).strip() != runner_task: + raise ValueError(f"Canonical and HF runner_task disagree for ESI-Bench question {question_id}.") + task["question_relpath"] = canonical_path.resolve().relative_to(canonical_root.resolve()).as_posix() + tasks.append(task) + return tasks + + +def load_datasets( + questions_jsonl: Path, + train_split: Path, + validation_split: Path, + *, + max_steps: int = 30, + canonical_root: Path | None = None, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Load explicit disjoint optimization and fixed-validation datasets.""" + + if max_steps < 1: + raise ValueError("max_steps must be positive.") + train_ids = read_split(train_split) + validation_ids = read_split(validation_split) + overlap = sorted(set(train_ids) & set(validation_ids)) + if overlap: + raise ValueError("ESI-Bench train/validation splits overlap: " + ", ".join(overlap)) + rows = load_question_rows(questions_jsonl) + canonical_questions = index_canonical_questions(canonical_root) if canonical_root is not None else None + return ( + materialize_split( + rows, + train_ids, + max_steps=max_steps, + canonical_questions=canonical_questions, + canonical_root=canonical_root, + ), + materialize_split( + rows, + validation_ids, + max_steps=max_steps, + canonical_questions=canonical_questions, + canonical_root=canonical_root, + ), + ) + + +def task_ids(tasks: Iterable[Mapping[str, Any]]) -> list[str]: + """Return stable task ids for split and provenance tests.""" + + return [str(task["task_id"]) for task in tasks] diff --git a/contrib/recipes/shaper/esi_bench/evaluate.py b/contrib/recipes/shaper/esi_bench/evaluate.py new file mode 100644 index 000000000..9b2d25380 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/evaluate.py @@ -0,0 +1,24 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Evaluate a SHAPER artifact pair on ESI-Bench.""" + +from __future__ import annotations + +from typing import Sequence + +from ..cli import cli_arguments, print_preflight_errors, requests_help +from ..evaluate import main as evaluate_main +from .train import FACTORY, preflight_environment + + +def main(argv: Sequence[str] | None = None) -> int: + arguments = cli_arguments(argv) + if requests_help(arguments): + return evaluate_main(["--factory", FACTORY, *arguments]) + if print_preflight_errors(preflight_environment()): + return 2 + return evaluate_main(["--factory", FACTORY, *arguments]) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/contrib/recipes/shaper/esi_bench/factory.py b/contrib/recipes/shaper/esi_bench/factory.py new file mode 100644 index 000000000..cd0973ee6 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/factory.py @@ -0,0 +1,180 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Environment-configured SHAPER bundle for official ESI-Bench rollouts.""" + +from __future__ import annotations + +import json +import os +import sys +from pathlib import Path +from typing import Any + +from agentlightning.types import LLM, PromptTemplate +from contrib.recipes.shaper.reproduce import ReproductionBundle + +from ..common import load_text +from .agent import ESIBenchAgent, ESIBenchRuntimeConfig +from .check_env import absolute_executable, check_map_generation_patch, check_worker_environment +from .contracts import ( + BEHAVIOR_ASSET_VERSION, + BEHAVIOR_COMMIT, + BEHAVIOR_REPOSITORY, + HARNESS_CONTRACT, + OMNIGIBSON_ROBOT_ASSET_VERSION, + UPSTREAM_COMMIT, + UPSTREAM_REPOSITORY, + check_behavior_source, + check_upstream_source, + make_harness_validator, + validate_skill, +) +from .dataset import load_datasets, task_ids +from .roles import ESIBenchRoleProtocol + +RECIPE_DIR = Path(__file__).parent +PROMPT_DIR = RECIPE_DIR / "prompts" +SPLIT_DIR = RECIPE_DIR / "splits" + + +def _required_env(name: str) -> str: + value = os.environ.get(name, "").strip() + if not value: + raise ValueError(f"Set {name} before building the ESI-Bench SHAPER bundle.") + return value + + +def _planner_resource() -> LLM: + api_key_env = os.environ.get("SHAPER_API_KEY_ENV", "OPENAI_API_KEY") + sampling: dict[str, Any] = { + "max_completion_tokens": int(os.environ.get("SHAPER_PLANNER_MAX_TOKENS", "32768")), + "optimizer_max_completion_tokens": int(os.environ.get("SHAPER_OPTIMIZER_MAX_TOKENS", "65536")), + "timeout": float(os.environ.get("SHAPER_PLANNER_TIMEOUT", "300")), + "max_retries": int(os.environ.get("SHAPER_PLANNER_RETRIES", "2")), + "temperature": float(os.environ.get("SHAPER_PLANNER_TEMPERATURE", "1.0")), + "top_p": float(os.environ.get("SHAPER_PLANNER_TOP_P", "0.95")), + "presence_penalty": float(os.environ.get("SHAPER_PLANNER_PRESENCE_PENALTY", "0.0")), + } + extra_body = os.environ.get("SHAPER_PLANNER_EXTRA_BODY", "").strip() + if extra_body: + parsed: object = json.loads(extra_body) + if not isinstance(parsed, dict): + raise ValueError("SHAPER_PLANNER_EXTRA_BODY must be a JSON object.") + sampling["extra_body"] = parsed + return LLM( + endpoint=_required_env("SHAPER_PLANNER_ENDPOINT"), + model=_required_env("SHAPER_MODEL"), + api_key=os.environ.get(api_key_env), + sampling_parameters=sampling, + ) + + +def build_bundle() -> ReproductionBundle[dict[str, Any]]: + """Build a complete official-runner ESI-Bench training bundle.""" + + root = Path(_required_env("ESI_BENCH_ROOT")).expanduser().resolve() + behavior_root = Path(_required_env("ESI_BEHAVIOR_ROOT")).expanduser().resolve() + worker_python = absolute_executable(Path(os.environ.get("ESI_WORKER_PYTHON", sys.executable))) + source_errors = check_upstream_source(root) + if source_errors: + raise RuntimeError("Unsupported ESI-Bench checkout: " + "; ".join(source_errors)) + behavior_errors = [*check_behavior_source(behavior_root), *check_worker_environment(worker_python, behavior_root)] + if behavior_errors: + raise RuntimeError("Unsupported BEHAVIOR/OmniGibson environment: " + "; ".join(behavior_errors)) + map_errors = check_map_generation_patch(Path(_required_env("ESI_MAKE_MAPS_PATH"))) + if map_errors: + raise RuntimeError("Unsupported OmniGibson map setup: " + "; ".join(map_errors)) + output_root = Path(os.environ.get("ESI_OUTPUT_ROOT", "outputs/shaper/esi_runner")).expanduser().resolve() + raw_data_root = os.environ.get("ESI_OMNIGIBSON_DATA_ROOT") or os.environ.get("OMNIGIBSON_DATA_PATH") + if not raw_data_root: + raise ValueError("Set ESI_OMNIGIBSON_DATA_ROOT or OMNIGIBSON_DATA_PATH before building the bundle.") + omnigibson_data_root = Path(raw_data_root).expanduser().resolve() + questions_jsonl = ( + Path(os.environ.get("ESI_QUESTIONS_JSONL", str(root / "hf_dataset" / "data" / "questions.jsonl"))) + .expanduser() + .resolve() + ) + runtime = ESIBenchRuntimeConfig( + esi_bench_root=root, + behavior_root=behavior_root, + questions_jsonl=questions_jsonl, + output_root=output_root, + omnigibson_data_root=omnigibson_data_root, + worker_python=worker_python, + max_steps=int(os.environ.get("ESI_MAX_STEPS", "30")), + min_steps=int(os.environ.get("ESI_MIN_STEPS", "3")), + confidence_threshold=float(os.environ.get("ESI_CONFIDENCE_THRESHOLD", "0.85")), + max_new_tokens=int(os.environ.get("ESI_MAX_NEW_TOKENS", "32768")), + temperature=float(os.environ.get("ESI_TEMPERATURE", "1.0")), + top_p=float(os.environ.get("ESI_TOP_P", "0.95")), + robot=os.environ.get("ESI_ROBOT", "R1"), + episode_timeout_seconds=float(os.environ.get("ESI_EPISODE_TIMEOUT", "1800")), + environment_retries=int(os.environ.get("ESI_ENVIRONMENT_RETRIES", "1")), + harness_timeout_seconds=float(os.environ.get("SHAPER_HARNESS_TIMEOUT", "3")), + harness_memory_limit_mb=int(os.environ.get("SHAPER_HARNESS_MEMORY_MB", "768")), + harness_max_output_chars=int(os.environ.get("SHAPER_HARNESS_MAX_OUTPUT_CHARS", "24000000")), + ) + train_split = Path(os.environ.get("ESI_TRAIN_SPLIT", str(SPLIT_DIR / "recipe_train10.txt"))).expanduser().resolve() + validation_split = ( + Path(os.environ.get("ESI_VALIDATION_SPLIT", str(SPLIT_DIR / "recipe_validation10.txt"))).expanduser().resolve() + ) + train, validation = load_datasets( + questions_jsonl, + train_split, + validation_split, + max_steps=runtime.max_steps, + canonical_root=root / "dataset" / "json_clean", + ) + validator = make_harness_validator( + timeout_seconds=runtime.harness_timeout_seconds, + memory_limit_mb=runtime.harness_memory_limit_mb, + max_output_chars=runtime.harness_max_output_chars, + ) + resources = { + "planner_llm": _planner_resource(), + "skill": PromptTemplate(template=load_text(PROMPT_DIR, "seed_skill.txt"), engine="f-string"), + "harness": PromptTemplate(template=load_text(PROMPT_DIR, "seed_harness.py"), engine="f-string"), + } + return ReproductionBundle( + agent=ESIBenchAgent(runtime), + train_dataset=train, + val_dataset=validation, + initial_resources=resources, + planner_resource_name="planner_llm", + harness_contract=HARNESS_CONTRACT, + skill_validator=validate_skill, + harness_validator=validator, + role_protocol=ESIBenchRoleProtocol(PROMPT_DIR), + provenance={ + "implementation_scope": "SHAPER method implementation with a benchmark-specific interface and prompt pack", + "benchmark": "ESI-Bench", + "upstream_repository": UPSTREAM_REPOSITORY, + "upstream_commit": UPSTREAM_COMMIT, + "behavior_repository": BEHAVIOR_REPOSITORY, + "behavior_commit": BEHAVIOR_COMMIT, + "behavior_asset_version": BEHAVIOR_ASSET_VERSION, + "omnigibson_robot_asset_version": OMNIGIBSON_ROBOT_ASSET_VERSION, + "split_status": ( + "deterministic contrib 10-question optimization and 10-question fixed-validation recipe; " + "not claimed to be the unavailable original experiment manifest" + ), + "train_task_ids": task_ids(train), + "validation_task_ids": task_ids(validation), + "runner": "official active_explore.pipeline.run_one in a fresh process per episode", + "worker_python": str(worker_python), + "harness_interception_scope": ( + "every frozen-planner user context passes through the selected harness: official primary " + "collect_contents calls and the audited inclined-plane post_action_query call" + ), + "official_auxiliary_model_calls": ( + "the pinned runner's sole task-specific call is inclined-plane post-action analysis; " + "its official prompt, frame order, schema, and per-call token limit remain authoritative" + ), + "reward": "official task scorer exact-match result", + "observable_context": ( + "official RGB/reference images, pixel-only derivatives, visible action/reasoning/confidence, " + "and sanitized official action results" + ), + "prompt_pack": "contrib/recipes/shaper/esi_bench/prompts", + }, + ) diff --git a/contrib/recipes/shaper/esi_bench/patches/behavior_floor_maps.patch b/contrib/recipes/shaper/esi_bench/patches/behavior_floor_maps.patch new file mode 100644 index 000000000..109a01b1d --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/patches/behavior_floor_maps.patch @@ -0,0 +1,12 @@ +diff --git a/asset_pipeline/b1k_pipeline/usd_conversion/make_maps.py b/asset_pipeline/b1k_pipeline/usd_conversion/make_maps.py +--- a/asset_pipeline/b1k_pipeline/usd_conversion/make_maps.py ++++ b/asset_pipeline/b1k_pipeline/usd_conversion/make_maps.py +@@ -19,7 +19,7 @@ WALL_CATEGORIES = ["walls", "rail_fence"] + FLOOR_CATEGORIES = ["floors", "driveway", "lawn"] + DOOR_CATEGORIES = ["door", "sliding_door", "garage_door", "gate"] + IGNORE_CATEGORIES = ["carpet"] +-NEEDED_STRUCTURE_CATEGORIES = FLOOR_CATEGORIES + WALL_CATEGORIES ++NEEDED_STRUCTURE_CATEGORIES = FLOOR_CATEGORIES + + # Segmentation maps will be generated with the data from the below map's overlap query + GENERATE_SEG_MAPS_DURING_FNAME = "floor_trav_no_obj_0.png" diff --git a/contrib/recipes/shaper/esi_bench/prompts/episode_summarizer.txt b/contrib/recipes/shaper/esi_bench/prompts/episode_summarizer.txt new file mode 100644 index 000000000..212ae71fe --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/prompts/episode_summarizer.txt @@ -0,0 +1,73 @@ +You are a batch summarizer for APCO optimization on ESI-Bench active +exploration. Your job is to analyze a candidate's complete development-batch +evaluation and produce a concise summary that helps an optimizer understand +what went well and what went wrong. + +Each ESI-Bench episode is one embodied question-answer trajectory. A diagnostic +judger has already analyzed every valid trajectory. You must summarize patterns +across questions rather than re-judge individual answers or prescribe changes. + +The official ESI-Bench exact-match results are authoritative. Preserve the +supplied numerator and denominator exactly. Evidence-progress scores are dense +diagnostic signals only and must never be presented as benchmark accuracy. + +## Input + +You will receive: +- Candidate identity +- Official exact-match result for the development batch +- Counts of valid and environment-invalid trajectories +- Category-level results +- Aggregate execution statistics +- Per-question diagnostic judger reports + +## Your job + +Produce one concise summary of approximately 180-300 words that identifies: + +1. **Outcome**: State the official exact-match result, the number of valid + trajectories, and the number of environment-invalid trajectories. Mention + notable category-level variation without treating a small category as a + global trend. + +2. **What the agent did**: Describe the dominant exploration behavior across + trajectories, including typical action patterns, step use, viewpoint changes, + interaction attempts, and when the agent submitted an answer. + +3. **Cross-question patterns**: Describe repeated progress or failure patterns, + such as informative exploration, repeated/inverse actions, viewpoint + deadlock, premature commitment, persistent uncertainty, invalid actions, or + exhaustion of the 30-step budget. Distinguish systematic patterns from + isolated events. + +4. **Evidence acquisition**: Summarize whether the agent usually found no + relevant evidence, located the target/anchor, acquired only partial evidence, + or obtained sufficient evidence but reasoned or answered incorrectly. Note + recurring failures involving official references, cross-view correspondence, + temporal evidence, count coverage, geometry, or interaction outcomes only + when supported by multiple reports. + +5. **Failure mode, if any**: State the best-supported root causes from the + judger reports. Infer what the evidence supports, but do not force one cause + when failures are mixed or ambiguous. + +6. **Context effectiveness**: Describe whether the supplied context preserved + useful references and historical observations, allowed decisive evidence to + coexist, clearly separated current from retained evidence, or instead became + redundant, stale, distracting, badly ordered, or incomplete. + +7. **Environment separation**: Report genuine environment-invalid trajectories + separately. Do not use them as evidence of agent, prompt, or context quality. + A wrong answer, model-generated invalid action, timeout at the 30-step budget, + or long trajectory is not an environment failure. + +## Output format + +Write a single paragraph. Do not use bullet points, headers, JSON, or tables. +Be specific and evidence-grounded. You may quote a recurring Reason/Action +pattern when useful, but do not list question IDs, scene names, or ground-truth +answers. + +Describe what happened rather than prescribing fixes. Do not recommend a prompt +edit, a context-code edit, an optimization target, or a replacement artifact. +Leave fix-finding to the optimizer. diff --git a/contrib/recipes/shaper/esi_bench/prompts/episode_summarizer_user.txt b/contrib/recipes/shaper/esi_bench/prompts/episode_summarizer_user.txt new file mode 100644 index 000000000..233fa1a4d --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/prompts/episode_summarizer_user.txt @@ -0,0 +1,20 @@ +Candidate: {candidate_name} + +Official exact-match result: +{n_correct}/{n_scored} = {score} + +Valid trajectories: {n_valid} +Environment-invalid trajectories: {n_environment_invalid} + +Category-level results: +{category_summary} + +Aggregate execution statistics: +{execution_statistics} + +Per-question diagnostic judger reports: +{judgements} + +Summarize the observed development-batch outcome and patterns. Preserve the +official score exactly, separate environment-invalid runs, and leave all +fix-finding to the optimizer. diff --git a/contrib/recipes/shaper/esi_bench/prompts/harness_optimizer.txt b/contrib/recipes/shaper/esi_bench/prompts/harness_optimizer.txt new file mode 100644 index 000000000..42faa59e3 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/prompts/harness_optimizer.txt @@ -0,0 +1,263 @@ +You are the context-harness optimizer for APCO on ESI-Bench active +exploration. + +Your job is to analyze development-batch execution feedback and produce one +improved context-builder Python module. The context builder controls which +official observations and trajectory records reach the embodied planner, how +they are labeled and ordered, and which deterministic pixel-derived views are +added. The planner skill is frozen and read-only. + +## System architecture + +1. **Official task prompt**: Generated by ESI-Bench for the current question. + It defines the question, answer options, task-specific evidence, legal + actions, and required Reason/Action interface. + +2. **Frozen planner skill**: Adds reusable evidence-acquisition and decision + policy around the official task prompt. It is supplied read-only and must + not be rewritten by this optimizer. + +3. **Context harness**: The only editable artifact. Before every planner step, + `build_context(...)` converts the official current observation, official + references, and observable execution history into an ordered multimodal + payload. + +4. **Embodied planner**: Receives the frozen skill and the payload emitted by + the harness, then returns exactly one legal exploration, interaction, or + answer action. + +5. **Environment**: Executes the action and returns the next official + observation or observable action result. An episode has at most 30 steps. + +The optimization schedule is fixed: the planner skill has already been +selected and is now frozen. Modify only the context harness. + +## Runtime interface + +The replacement module must define: + +```python +def build_context( + image_path: Path, + history: list[dict[str, Any]], + prompt: str, + reference_image_paths: list[Path] | None = None, +) -> list[Any]: +``` + +`build_context` is called before each planner action. + +- `image_path` is the current official egocentric RGB observation. +- `reference_image_paths` contains zero or more official question-reference + images. +- `prompt` contains the current official question, options, legal actions, and + response contract. The question and action contract are authoritative. +- `history` contains records from already executed steps. Observable fields + include `step`, `action`, `answer`, `confidence`, `reasoning`, `image`, + `extra_image_paths`, and `action_result`. +- A history `image` is an official prior RGB filename relative to + `image_path.parent`. +- `extra_image_paths` may contain official action-generated past views. +- `reasoning` is fallible model-written text. It may summarize what the model + believed it saw, but it is not ground truth. +- Only visibly returned action-result fields may be used, such as `handled`, + `operation`, `action`, `success`, `error`, `reason`, `object`, `target`, + `container`, `physical_state`, `attempts`, and `current_stack`. + +The returned list is sent to the planner in order. Each item must be either: + +- a string text block, or +- a `pathlib.Path` to an official image or a deterministic image derived only + from official RGB pixels. + +Ordering and labels are part of the artifact. The planner must be able to +distinguish official references, retained historical views, derived crops or +overlays, action-generated extra views, the current view, the official task +contract, and harness directives. + +## Observable-information boundary + +The harness may use only: + +- Official question-reference images +- Current and historical official RGB observations +- Official action-generated extra views +- Executed actions and their visibly returned results +- Planner reasoning, answers, and confidence as fallible records +- Question text, answer options, legal action vocabulary, and output contract +- Deterministic calculations over the items above + +The harness must never read or infer from simulator object poses, camera pose, +object metadata, scale, volume, depth, segmentation, semantic maps, collision +state, hidden phase state, ground-truth answers, or any other privileged +simulator channel. A `camera` or metadata field must be ignored even if it is +present in a raw runner record. + +Action counts may be integrated into explicitly labeled relative commanded +motion, but must not be represented as measured simulator pose. + +## Clean pixel transformations + +The harness may generate deterministic derivatives of official RGB images when +they expose existing pixels without inventing observations. Examples include: + +- Bounded focus crops selected from fixed candidate windows +- Pixel-quality measurements such as contrast, edge density, and Laplacian + sharpness +- Low-information and visual-stagnation tests on downsampled RGB +- Deterministic coordinate grids or annotations over an official frame +- Aspect-correct arithmetic over pixel coordinates reported by the planner + +Every derived image must be labeled as a crop or overlay of a named official +frame. A crop is not a new viewpoint and must not be presented as one. Derived +files may be cached under the current episode image directory using stable +names. The module must degrade gracefully to the original official image when +decoding or writing a derivative fails. + +For a frozen planner that uses normalized visual-grounding coordinates, a +matching deterministic grid is allowed. When arithmetic is performed, restore +the official image aspect ratio before measuring distances or offsets. The +grid and arithmetic may use only official pixels and planner-reported points; +they may not obtain coordinates from simulator state. + +## Context-harness design space + +The optimizer should discover the smallest systematic changes supported by the +development evidence. Useful changes can include, but are not limited to: + +### Text trajectory management + +- Replace an unbounded text dump with a compact action path, recent observable + results, and recent model observations. +- Mark model-written observations as fallible. +- Preserve facts needed to interpret retained frames while removing duplicated + runner history that competes with current evidence. +- Keep the official question, options, legal actions, and response contract + intact. + +### Sparse visual memory + +- Retain a small bounded set of evidence-bearing historical frames rather than + only the most recent fixed window. +- Consider an informative initial anchor, the strongest task-relevant pixel + evidence, a recent informative view, and visually diverse evidence. +- Reject blank, obstructed, redundant, or low-information frames. +- Combine RGB quality, task relevance, recency, and visual diversity without + treating model text as ground truth. +- Include both the retained full frame and a bounded deterministic focus crop + when the crop exposes small but already visible evidence. + +### Current evidence + +- Always label the current official full frame unambiguously. +- Add a bounded number of current-view focus crops only when useful. +- Keep current and retained observations jointly available when cross-view + comparison is required. +- Preserve official reference images and distinguish them from agent-observed + views. + +### Evidence and recovery control + +- Detect repeated actions, inverse-action cycles, low-information views, or + visually stagnant transitions from observable history. +- Inject a bounded recovery directive only when its trigger is present. +- Track the remaining step budget and require a terminal answer when + exploration closes. +- Stop a camera macro as soon as the required evidence is visibly ready. +- Preserve the key evidence frame used at a readiness transition so it remains + available on the following answer step. + +### Task-conditional evidence routing + +- Parse a broad task family from the official question when repeated feedback + supports a reusable family-level policy. +- Cross-view tasks may need a reflected, occluded, or earlier anchor jointly + routed with current candidates. +- Counting may need coverage-aware retained views and duplicate avoidance. +- Temporal tasks may need official references and observations from distinct + phases kept together. +- Geometry tasks may need a bounded elevation/pitch acquisition, an official + RGB coordinate overlay, explicit target-point reporting, a retained + readiness bookmark, and deterministic pixel arithmetic. +- Physical interaction tasks may need compact observable action-result state + and recovery after a legal interaction fails. + +These are allowed mechanisms, not a mandatory template. Add a mechanism only +when summaries, reports, payload traces, or validation feedback justify it. +Prefer bounded general policies over a catalogue of per-question recipes. + +## Inputs and diagnosis + +You receive: + +- The frozen planner skill +- The current context-builder module +- Official development accuracy and validity counts +- A neutral batch summary +- Representative official task interfaces +- Representative diagnostic Judger reports +- Context-execution traces showing the actual ordered payload emitted at each + selected step +- Representative planner/action traces and outcomes +- Artifact validation or smoke-test feedback +- Optimization history + +Use payload traces to reason about what the planner actually received, not what +the source code was intended to provide. Diagnose whether failures arise from +missing evidence, harmful context, bad ordering or labels, excessive context, +stale memory, poor crop selection, loop handling, readiness/stop logic, +planner reasoning despite sufficient payload, or a true environment failure. + +Do not modify the harness to compensate for a planner-only reasoning error +unless a concrete payload change could expose or organize relevant official +evidence more effectively. Do not optimize against environment-invalid runs as +if they were ordinary wrong answers. + +## Implementation contract + +- Return one complete replacement Python module. +- Define `build_context` with the exact signature above. +- Do not mutate `history` or any history record. +- The module must be deterministic and have bounded per-step work. +- Standard-library modules plus `pathlib`, `typing`, `numpy`, and `cv2` may be + used. Local helper modules already supplied with the current candidate may be + imported or loaded read-only; do not invent unavailable dependencies or load + code from outside the supplied context-builder artifact directory. +- Reading official image paths and writing deterministic derived images beneath + `image_path.parent` is allowed. Arbitrary file access is forbidden. +- Network access, subprocesses, generated-code execution, `eval`, `exec`, and + access to process or global secrets are forbidden. +- Cache image decoding or derivatives by stable path when useful, but ensure a + cold process still behaves correctly. +- Catch expected image/path failures and fall back to a valid payload rather + than crashing. +- Keep payload size bounded. More images are not automatically better. +- Preserve the frozen skill's Reason/Action fields and the official legal + action vocabulary. + +## Generalization constraints + +- Do not hard-code question IDs, scene names, development examples, answer + positions, ground-truth answers, or benchmark lookup tables. +- Do not select frames using correctness labels or information unavailable at + evaluation time. +- Task-family routing must be reusable and derived from the official question + or observable trajectory. +- Never claim that a crop, overlay, model observation, or action-derived + estimate is privileged ground truth. +- Preserve behavior that already works unless the feedback supports changing + it. +- Prefer the smallest coherent improvement with explicit regression risks. +- Do not output a new planner skill, a diff, commentary around the artifact, or + multiple candidates. + +## Output + +Return exactly one JSON object: + +{ + "rationale": "", + "new_artifact": "" +} + +No code fences and no surrounding prose. diff --git a/contrib/recipes/shaper/esi_bench/prompts/harness_optimizer_user.txt b/contrib/recipes/shaper/esi_bench/prompts/harness_optimizer_user.txt new file mode 100644 index 000000000..b0cbe79c7 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/prompts/harness_optimizer_user.txt @@ -0,0 +1,67 @@ +FROZEN PLANNER SKILL (READ-ONLY) +================================ +{frozen_planner_skill} + +CURRENT CONTEXT-HARNESS MODULE +============================== +{current_context_code} + +OFFICIAL DEVELOPMENT RESULT +=========================== +Candidate: {candidate_name} +Official exact-match score: {n_correct}/{n_scored} = {score} +Valid trajectories: {n_valid} +Environment-invalid trajectories: {n_environment_invalid} + +NEUTRAL BATCH SUMMARY +===================== +{batch_summary} + +REPRESENTATIVE OFFICIAL TASK INTERFACES +======================================= +The following anonymized excerpts contain task questions, reference-image +availability, legal action vocabularies, and output contracts, but no +ground-truth answers: + +{representative_task_interfaces} + +REPRESENTATIVE DIAGNOSTIC JUDGER REPORTS +======================================== +{representative_judger_reports} + +REPRESENTATIVE CONTEXT-EXECUTION TRACES +======================================= +Each excerpt describes the observable input history and the actual ordered +text/image payload emitted by the current harness at selected planner steps. +Image entries include source step, dimensions, derivative lineage, and compact +pixel statistics when available: + +{representative_context_payloads} + +REPRESENTATIVE PLANNER/ACTION TRACE EXCERPTS +============================================ +{representative_trace_excerpts} + +ARTIFACT VALIDATION AND SMOKE-TEST FEEDBACK +=========================================== +{artifact_validation_feedback} + +OPTIMIZATION HISTORY +==================== +{optimization_history} + +Produce one improved context-harness module for the next candidate. The +planner skill is frozen and read-only. Address systematic payload failures +without hard-coding development questions or using privileged simulator +information. + +The returned `new_artifact` must expose an interface compatible with: + +def build_context(image_path, history, prompt, reference_image_paths=None) -> list + +Return exactly one JSON object: + +{{ + "rationale": "", + "new_artifact": "" +}} diff --git a/contrib/recipes/shaper/esi_bench/prompts/round_judger.txt b/contrib/recipes/shaper/esi_bench/prompts/round_judger.txt new file mode 100644 index 000000000..79346bd82 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/prompts/round_judger.txt @@ -0,0 +1,163 @@ +You are the diagnostic judger for APCO on ESI-Bench active exploration. + +Each ESI-Bench episode is an embodied question-answer trajectory. The agent +may move the camera or manipulate the scene for up to 30 steps before +submitting an answer. Analyze the complete trajectory, not an isolated action. + +The official ESI-Bench exact-match scorer supplies the binary correctness +label. That label is authoritative. You must copy it exactly and must not +replace it with your own visual judgement. Your evidence-progress score is a +dense diagnostic signal for optimization only; it is never a reported +benchmark score. + +## Input + +You receive: +- Task family, scene, question, answer options, and official task contract +- Final answer, confidence, termination reason, and official exact-match result +- The official ground truth, available only for post-hoc development diagnosis +- Per-step Reason/Action history, action validity, and action results +- Official question-reference images, when provided by the benchmark +- An audit view of evidence-bearing historical observations and final views +- The planner prompt augmentation and the actual context supplied at the final + decision +- Environment diagnostics, including initialization, renderer, or runner faults + +Use only the supplied official observations and execution records. Do not +assume access to simulator poses, object metadata, depth, segmentation, hidden +state, or any other privileged information. + +## Your role + +Objectively determine: +1. Whether the run was valid and whether the official answer was correct. +2. What evidence the question required and what the agent actually observed. +3. Whether exploration acquired new information or became inefficient/stuck. +4. Whether the final answer and confidence were supported by the evidence. +5. Whether the planner prompt or context harness caused the observed behavior. +6. What single generalizable change would most help the optimizer. + +## Evaluation criteria + +### 1. Official outcome and run validity + +- Copy `official_correct` exactly from the official scorer. +- Set `run_valid=false` only when an environment, renderer, or runner failure + prevented a meaningful trajectory. +- A wrong answer, max-step termination, long trajectory, illegal model action, + repeated action, or failure to answer is an agent failure, not an environment + failure. +- When `run_valid=false`, set `evidence_progress_score` to null and attribute + the failure to `environment`. + +### 2. Evidence-progress score + +For a valid run, select the highest stage clearly supported by the supplied +images, references, trajectory, and context. Output exactly one of +{0.0, 0.25, 0.5, 0.75, 1.0}: + +- **0.00 - No relevant evidence**: The agent did not locate a task-relevant + object, region, relation, event, or reference anchor. +- **0.25 - Relevant target located**: The agent found a relevant object, room, + region, or reference anchor, but did not obtain enough evidence to compare or + answer. +- **0.50 - Partial evidence acquired**: The trajectory contains useful evidence, + but a required viewpoint, cross-view correspondence, temporal observation, + count coverage, or physical outcome is still missing. +- **0.75 - Sufficient evidence acquired**: The observations appear sufficient + to answer, but the agent reasoned incorrectly, used the evidence poorly, + failed to submit, or submitted an unsupported answer. +- **1.00 - Correctly completed**: The official exact-match scorer marks the + final answer correct. + +`evidence_progress_score` must be 1.0 when `official_correct=true`. It must not +be used to alter the official binary result. + +### 3. Evidence acquisition + +- Identify the minimum visual or interaction evidence required by the question. +- State which evidence was observed and at which steps. +- Distinguish current-view evidence, official reference evidence, retained + cross-view memory, and interaction outcomes. +- Check whether decisive evidence existed earlier but disappeared from the + actual context before the final decision. +- Do not infer that evidence exists merely because it would support the ground + truth; cite supplied observations or records. + +### 4. Exploration quality + +- Check action legality against the official task contract. +- Determine whether each action sought a specific missing or falsifying piece + of evidence. +- Detect repeated or inverse action cycles, blank/cropped viewpoints, motion + without new evidence, viewpoint deadlock, and failure to recover perspective. +- Diagnose both premature commitment and unnecessary exploration after enough + evidence was available. +- Treat exhaustion of the 30-step budget as an agent outcome unless an + infrastructure fault caused the delay. + +### 5. Answer and confidence + +- Check whether the final answer is supported by observations available to the + planner at the decision step. +- Diagnose premature commitment, unjustified high confidence, persistent + uncertainty despite sufficient evidence, or failure to answer before budget. +- The ground truth may explain a development failure, but never recommend a + question ID, scene name, answer key, or task-specific lookup rule. + +### 6. Prompt and context effectiveness + +Choose `optimizer_target="prompt"` when the dominant problem is: +- Task or official-contract misunderstanding +- Poor evidence-seeking or falsification policy +- Illegal or badly formatted actions +- Premature commitment or confidence discipline +- Incorrect reasoning despite sufficient supplied evidence + +Choose `optimizer_target="context"` when the dominant problem is: +- Required evidence was observed but omitted or evicted +- Cross-view or temporal evidence was not retained together +- References, observations, crops, labels, or history were badly selected, + ordered, compressed, repeated, or presented too late +- The planner could not distinguish current evidence from stale evidence + +Choose `optimizer_target="environment"` only for a genuine invalid run, and +choose `optimizer_target="none"` when the run is correct and no systematic +defect is visible. + +## Primary failure taxonomy + +Use exactly one label: +- none +- missing_evidence +- lost_cross_view_evidence +- reference_ignored +- action_blindness +- invalid_action +- viewpoint_deadlock +- premature_commitment +- reasoning_error +- confidence_miscalibration +- budget_exhaustion +- task_contract_violation +- environment_failure +- other + +## Output + +Return one JSON object only: + +{ + "official_correct": true, + "run_valid": true, + "evidence_progress_score": 1.0, + "evidence_analysis": "", + "exploration_analysis": "", + "answer_analysis": "", + "context_analysis": "", + "primary_failure": "", + "optimizer_target": "prompt | context | environment | none", + "improvement_signal": "" +} + +Be concise, objective, evidence-grounded, and explicit about uncertainty. diff --git a/contrib/recipes/shaper/esi_bench/prompts/round_judger_user.txt b/contrib/recipes/shaper/esi_bench/prompts/round_judger_user.txt new file mode 100644 index 000000000..0f2a42745 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/prompts/round_judger_user.txt @@ -0,0 +1,39 @@ +Task family: {task_family} +Scene: {scene} +Question ID: {question_id} + +Question: +{question} + +Answer options: +{options} + +Official task contract: +{task_contract} + +Final answer: {final_answer} +Final confidence: {final_confidence} +Termination: {termination} +Official ground truth: {ground_truth} +Official exact-match result: {official_correct} + +Environment diagnostics: +{environment_diagnostics} + +Planner prompt augmentation: +{planner_prompt} + +Actual context supplied at the final decision: +{final_context} + +Per-step Reason/Action trajectory: +{trajectory} + +Images follow in this order: +1. Official question-reference images, if any +2. Evidence-bearing historical observations from the trajectory audit +3. Final/current observations + +Diagnose the complete trajectory without changing the official correctness +label. Use the ground truth only for post-hoc development diagnosis, never to +recommend a question-specific rule. diff --git a/contrib/recipes/shaper/esi_bench/prompts/seed_harness.py b/contrib/recipes/shaper/esi_bench/prompts/seed_harness.py new file mode 100644 index 000000000..2dd90eff5 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/prompts/seed_harness.py @@ -0,0 +1,64 @@ +# Copyright (c) Microsoft. All rights reserved. + + +def build_context(records): + past = [record for record in records if record.get("record_kind") == "past"] + current = next( + (record for record in reversed(records) if record.get("record_kind") == "current"), + {}, + ) + parts = [] + if current.get("call_kind") == "auxiliary_post_action": + for item in current.get("observable_sequence", []): + if item.get("content_kind") == "text": + parts.append({"type": "text", "text": str(item.get("text", ""))}) + elif item.get("content_kind") == "observation": + observation = item.get("observation", {}) + if isinstance(observation.get("full_frame"), dict): + parts.append(observation["full_frame"]) + elif item.get("content_kind") == "content_part" and isinstance(item.get("part"), dict): + parts.append(item["part"]) + return parts + if past: + lines = [] + for record in past: + lines.append( + "Step " + + str(record.get("step", 0)) + + ": action=" + + str(record.get("action", "")) + + " answer=" + + str(record.get("answer", "")) + + " confidence=" + + str(record.get("confidence", 0.0)) + + " result=" + + str(record.get("action_result_text", "{}")) + + " reasoning=" + + str(record.get("reasoning", "")) + ) + parts.append({"type": "text", "text": "Action history so far:\n" + "\n".join(lines)}) + for reference in current.get("reference_observations", []): + parts.append({"type": "text", "text": str(reference.get("label", "QUESTION REFERENCE IMAGE"))}) + if isinstance(reference.get("full_frame"), dict): + parts.append(reference["full_frame"]) + for record in past[-5:]: + observation = record.get("observation", {}) + parts.append({"type": "text", "text": "[Past view from step " + str(record.get("step", 0)) + "]"}) + if isinstance(observation.get("full_frame"), dict): + parts.append(observation["full_frame"]) + for extra_index, extra in enumerate(record.get("extra_observations", []), start=1): + parts.append( + { + "type": "text", + "text": ( + "[Past extra view from step " + str(record.get("step", 0)) + " #" + str(extra_index) + "]" + ), + } + ) + if isinstance(extra.get("full_frame"), dict): + parts.append(extra["full_frame"]) + observation = current.get("observation", {}) + parts.append({"type": "text", "text": "[CURRENT VIEW - step " + str(current.get("step", 1)) + "]"}) + if isinstance(observation.get("full_frame"), dict): + parts.append(observation["full_frame"]) + return parts diff --git a/contrib/recipes/shaper/esi_bench/prompts/seed_skill.txt b/contrib/recipes/shaper/esi_bench/prompts/seed_skill.txt new file mode 100644 index 000000000..8dd1d6e90 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/prompts/seed_skill.txt @@ -0,0 +1 @@ +{task_prompt} diff --git a/contrib/recipes/shaper/esi_bench/prompts/skill_optimizer.txt b/contrib/recipes/shaper/esi_bench/prompts/skill_optimizer.txt new file mode 100644 index 000000000..06fb31073 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/prompts/skill_optimizer.txt @@ -0,0 +1,162 @@ +You are the planner-skill optimizer for APCO on ESI-Bench active exploration. + +Your job is to analyze development-batch execution feedback and produce one +improved reusable planner skill prompt. The skill controls how the embodied +planner acquires evidence, uses multimodal context, chooses one next action, +and decides when to answer. + +## System architecture + +1. **Official task prompt**: Generated by ESI-Bench for the current question. + It defines the question, answer options, available evidence, legal action + vocabulary, task-specific contract, and required Reason/Action interface. + +2. **Planner skill**: A reusable prompt augmentation wrapped around the official + task prompt. This is the only artifact you may modify. + +3. **Frozen context harness**: Builds the ordered multimodal payload from + official question references, historical observations, action/reasoning + records, action results, and the current observation. Its complete source is + supplied read-only. You may use it to understand what the planner receives, + but you must not output or modify context code. + +4. **Embodied planner**: Receives the official task prompt, planner skill, frozen + harness payload, and current official observation. At every step it returns + exactly one legal exploration, manipulation, or answer action. + +5. **Environment**: Executes that action and returns the next official + observation or interaction result. An episode has a maximum of 30 steps. + +The optimization schedule is fixed: planner skill is optimized first while the +context harness is frozen. Do not choose an optimization target and do not try +to edit the harness indirectly. + +## Runtime planner interface + +The literal placeholder `{task_prompt}` represents the complete official +per-question prompt and must be preserved exactly once in every candidate skill. +Depending on the task, the runtime interface may provide: + +- Question text and listed answer options +- Official question-reference images +- Current egocentric RGB observation +- Prior Reason/Action history and observable action results +- A task-specific legal action vocabulary +- `answer(answer, confidence)` as the terminal action +- A required response form beginning with `Reason:` and `Action:` + +The official task contract is authoritative. A skill may add reusable reasoning +and evidence-acquisition policy, but must not replace, contradict, or duplicate +the per-question action schema. + +## ESI-Bench task and evidence distribution + +ESI-Bench contains heterogeneous embodied questions. Useful planner policy must +generalize across evidence types rather than assume that every question can be +solved from the current frame. + +- **Reference-grounded perception**: Some tasks provide authoritative reference + images or maps that must be compared with current or historical observations. +- **Cross-view perception**: Reflection, occlusion, spatial, and mapping tasks + may require evidence collected from different viewpoints to be considered + together. +- **Geometric and metric reasoning**: Distance, size, line, triangle, touching, + and related tasks may require an informative viewpoint before comparison. +- **Perceptual grounding**: Angle, occlusion, transparency, material, or object + identity questions require grounding conclusions in visible pixels rather + than object-list wording alone. +- **Physical interaction and dynamics**: Deformable, slope, stacking, storage, + pouring, and action-sequencing tasks may require a legal interaction followed + by observation of its outcome. +- **Temporal evidence**: Change and multi-agent tasks may require comparing + observations or outcomes across phases or time. +- **Enumerative evidence**: Counting requires sufficient coverage and a method + that avoids both omission and double counting. + +This distribution is not an answer key. Use the supplied summaries and +representative reports to determine which reusable policies are justified. + +## Common planner-skill failure modes + +- **Missing evidence**: The agent answers without locating the evidence required + by the question. +- **Reference ignored**: An official reference is available but not used. +- **Action blindness**: The agent repeats one action class without changing the + evidence state. +- **Viewpoint deadlock**: Repeated/inverse movements, blank views, cropped + targets, or poor perspective prevent informative observation. +- **Invalid action**: The planner violates the current task's legal action + vocabulary or output syntax. +- **Premature commitment**: The planner answers before checking the evidence + needed to distinguish the listed options. +- **Persistent uncertainty**: Sufficient evidence appears available, but the + planner continues exploring or ends with no supported answer. +- **Reasoning error**: The necessary evidence reaches the planner, but it is + interpreted or compared incorrectly. +- **Confidence miscalibration**: Confidence is inconsistent with the evidence. +- **Budget misuse**: The 30-step budget is spent without acquiring information + that could change the answer. + +## Your task + +Analyze the current skill, frozen harness, official development result, neutral +batch summary, representative diagnostic reports, representative task +interfaces, Reason/Action trace excerpts, and optimization history. + +Identify the smallest set of systematic planner-policy changes that addresses +the strongest repeated failures while preserving behavior that already works. +Every proposed rule must be supported by the supplied feedback or by the +official interface contract. + +The improved skill may introduce: + +- A compact evidence inventory or answer-readiness check +- Explicit use of official references and retained historical evidence +- A policy for seeking new or falsifying evidence +- Recovery from repeated actions, low-information views, or poor perspective +- Task-family-conditional evidence checks supported by repeated feedback +- Legal-action and exact-output discipline +- Calibrated confidence and timely answer submission + +These are design possibilities, not mandatory sections. Do not add a rule when +the supplied evidence does not justify it. Prefer a short coherent policy over +a catalogue of unrelated task recipes. + +## Component boundary + +- Modify only the planner skill prompt. +- Treat the supplied context harness as read-only. +- Do not compensate for absent context by claiming that the planner can see an + image, crop, label, metadata field, or history item that the harness does not + actually provide. +- If a failure is fundamentally caused by missing context, describe that limit + in the rationale but do not emit context code or fabricate a prompt-only fix. +- Preserve successful existing behavior unless feedback supports changing it. + +## Generalization and safety constraints + +- The complete improved skill must contain `{task_prompt}` exactly once. +- Preserve the official action vocabulary and Reason/Action response contract. +- Do not hard-code question IDs, scene names, development examples, option + positions, ground-truth answers, or benchmark-specific lookup tables. +- Do not request or use simulator poses, object metadata, depth, segmentation, + hidden state, privileged geometry, or any information outside the supplied + official observations and records. +- Task-family-conditional strategy is allowed only when it is reusable and + supported by repeated evidence; question-specific memorization is forbidden. +- Do not expose hidden chain-of-thought. The planner may be asked for a brief, + evidence-grounded `Reason:` field suitable for the benchmark interface. +- Do not output Python code, a diff, commentary around the artifact, or multiple + candidate prompts. +- Keep the improved skill concise, preferably no more than 100 lines. + +## Output + +Return exactly one JSON object: + +{ + "rationale": "", + "new_artifact": "" +} + +No code fences and no surrounding prose. diff --git a/contrib/recipes/shaper/esi_bench/prompts/skill_optimizer_user.txt b/contrib/recipes/shaper/esi_bench/prompts/skill_optimizer_user.txt new file mode 100644 index 000000000..007a273b3 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/prompts/skill_optimizer_user.txt @@ -0,0 +1,52 @@ +CURRENT PLANNER SKILL +===================== +{current_skill} + +FROZEN CONTEXT HARNESS (READ-ONLY) +================================== +{frozen_harness} + +OFFICIAL DEVELOPMENT RESULT +=========================== +Candidate: {candidate_name} +Official exact-match score: {n_correct}/{n_scored} = {score} +Valid trajectories: {n_valid} +Environment-invalid trajectories: {n_environment_invalid} + +NEUTRAL BATCH SUMMARY +===================== +{batch_summary} + +REPRESENTATIVE TASK INTERFACES +============================== +The following anonymized excerpts contain task contracts, evidence interfaces, +and legal action spaces but no ground-truth answers: + +{representative_task_interfaces} + +REPRESENTATIVE DIAGNOSTIC JUDGER REPORTS +======================================== +{representative_judger_reports} + +REPRESENTATIVE REASON/ACTION TRACE EXCERPTS +=========================================== +{representative_trace_excerpts} + +OPTIMIZATION HISTORY +==================== +{optimization_history} + +Produce one improved planner skill for the next candidate. The context harness +is frozen and read-only. Address systematic planner-policy failures without +hard-coding development questions or assuming unavailable evidence. + +The returned `new_artifact` must preserve the literal placeholder +{{task_prompt}} exactly once and must remain compatible with the official +Reason/Action interface. + +Return exactly one JSON object: + +{{ + "rationale": "", + "new_artifact": "" +}} diff --git a/contrib/recipes/shaper/esi_bench/roles.py b/contrib/recipes/shaper/esi_bench/roles.py new file mode 100644 index 000000000..f6b926061 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/roles.py @@ -0,0 +1,605 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Paper-faithful SHAPER role protocol for ESI-Bench.""" + +from __future__ import annotations + +import asyncio +import json +from collections import Counter, defaultdict +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence, cast + +from contrib.agentlightning.contrib.shaper import ( + ArtifactProposal, + CandidateEvaluation, + EpisodeTrace, + OptimizationStage, + OptimizerRequestContext, + RoleCompleter, + RoleRequest, + RoundRecord, + parse_json_object, +) + +from ..common import load_text + +_FAILURE_LABELS = { + "none", + "missing_evidence", + "lost_cross_view_evidence", + "reference_ignored", + "action_blindness", + "invalid_action", + "viewpoint_deadlock", + "premature_commitment", + "reasoning_error", + "confidence_miscalibration", + "budget_exhaustion", + "task_contract_violation", + "environment_failure", + "other", +} + +_AGL_HARNESS_ADAPTER = """ +## Agent Lightning execution adapter (authoritative for this contrib) + +The paper prompt above describes the original in-process ESI-Bench context +module. This contrib executes optimizer-generated code in an isolated JSON-only +worker. Preserve the same evidence-selection objective, but obey this concrete +interface: + +- Define exactly `def build_context(records)`. +- `records` contains only the observable, JSON-serializable fields documented + in the validation contract supplied in the user message. +- Official and derived images are already OpenAI `image_url` parts nested in + each observation. Reuse those parts directly; do not read or write paths. +- No `Path`, `numpy`, `cv2`, network client, process state, or simulator object + is available. Imports and external file access are rejected. +- Return a string or a bounded list of OpenAI `text` / `image_url` parts. +- The validation contract supplied in the user message overrides incompatible + low-level runtime details in the original prompt. +""".strip() + + +@dataclass(frozen=True) +class _JudgedTrace: + trace: EpisodeTrace + judgement: dict[str, Any] + + +@dataclass(frozen=True) +class _ESIBenchGradient: + evaluation: CandidateEvaluation + judged: tuple[_JudgedTrace, ...] + batch_summary: str + n_correct: int + n_scored: int + n_invalid: int + category_summary: str + execution_statistics: str + + +def _split_images(value: Any) -> tuple[Any, list[dict[str, Any]]]: + images: list[dict[str, Any]] = [] + + def visit(item: Any) -> Any: + if isinstance(item, list): + return [visit(child) for child in cast(list[Any], item)] + if not isinstance(item, dict): + return item + mapping = cast(dict[str, Any], item) + if mapping.get("type") == "image_url" and isinstance(mapping.get("image_url"), dict): + image = cast(dict[str, Any], mapping["image_url"]) + if isinstance(image.get("url"), str): + images.append(mapping) + return {"type": "image_url", "image_url": {"url": ""}} + return {str(key): visit(child) for key, child in mapping.items()} + + return visit(value), images + + +def _image_parts(parts: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]: + return [ + cast(dict[str, Any], part) + for part in parts + if part.get("type") == "image_url" and isinstance(part.get("image_url"), dict) + ] + + +def _metadata(trace: EpisodeTrace) -> dict[str, Any]: + return trace.metadata.extra + + +def _official_correct(trace: EpisodeTrace) -> bool: + value = _metadata(trace).get("official_correct") + return bool(value) if isinstance(value, bool) else float(trace.final_reward or 0.0) >= 1.0 + + +def _family(trace: EpisodeTrace) -> str: + value = str(_metadata(trace).get("task_family", "")).strip() + return value or "Unknown" + + +def _json_text(value: Any) -> str: + return json.dumps(value, ensure_ascii=False, indent=2, default=str) + + +def _selected_round_indices(rounds: Sequence[RoundRecord], limit: int = 6) -> list[int]: + if len(rounds) <= limit: + return list(range(len(rounds))) + candidates = [0, 1, len(rounds) // 2, len(rounds) - 2, len(rounds) - 1] + for index in range(1, len(rounds)): + if rounds[index].command != rounds[index - 1].command: + candidates.append(index) + output: list[int] = [] + for index in candidates: + if 0 <= index < len(rounds) and index not in output: + output.append(index) + if len(output) >= limit: + break + return sorted(output) + + +def _trajectory(trace: EpisodeTrace) -> str: + lines: list[str] = [] + for record in trace.rounds: + result = record.action_result + lines.append( + "Step " + + str(record.round_index + 1) + + "\nReason/Action output: " + + record.planner_response + + "\nParsed action: " + + record.command + + "\nAction valid: " + + str(result.get("action_valid", True)) + + "\nAction result: " + + _json_text(result) + ) + return "\n\n".join(lines) if lines else "(No executable planner step was recorded.)" + + +def _reference_images(trace: EpisodeTrace) -> list[dict[str, Any]]: + raw = _metadata(trace).get("reference_images") + if not isinstance(raw, list): + return [] + output: list[dict[str, Any]] = [] + for item in cast(list[Any], raw): + if not isinstance(item, dict): + continue + mapping = cast(dict[str, Any], item) + image = mapping.get("image") + if isinstance(image, dict): + image_mapping = cast(dict[str, Any], image) + if image_mapping.get("type") == "image_url": + output.append(image_mapping) + return output + + +def _deduplicate_images(images: Iterable[dict[str, Any]]) -> list[dict[str, Any]]: + output: list[dict[str, Any]] = [] + seen: set[str] = set() + for image in images: + raw = image.get("image_url") + url = str(cast(dict[str, Any], raw).get("url", "")) if isinstance(raw, dict) else "" + if url and url not in seen: + seen.add(url) + output.append(image) + return output + + +class ESIBenchRoleProtocol: + """Use whole-trajectory diagnosis and development-batch summarization.""" + + def __init__(self, prompt_dir: Path) -> None: + self.judger_system = load_text(prompt_dir, "round_judger.txt") + self.judger_user = load_text(prompt_dir, "round_judger_user.txt") + self.summarizer_system = load_text(prompt_dir, "episode_summarizer.txt") + self.summarizer_user = load_text(prompt_dir, "episode_summarizer_user.txt") + self.skill_optimizer_system = load_text(prompt_dir, "skill_optimizer.txt") + self.skill_optimizer_user = load_text(prompt_dir, "skill_optimizer_user.txt") + self.harness_optimizer_system = load_text(prompt_dir, "harness_optimizer.txt") + self.harness_optimizer_user = load_text(prompt_dir, "harness_optimizer_user.txt") + + @staticmethod + def _judger_images(trace: EpisodeTrace) -> list[dict[str, Any]]: + historical: list[dict[str, Any]] = [] + for index in _selected_round_indices(trace.rounds): + historical.extend(_image_parts(trace.rounds[index].observation_before)) + if trace.rounds: + _, context_images = _split_images(trace.rounds[-1].context_payload) + historical.extend(context_images) + final = _image_parts(trace.rounds[-1].observation_after) + else: + final = [] + content: list[dict[str, Any]] = [] + references = _reference_images(trace) + if references: + content.append({"type": "text", "text": "OFFICIAL QUESTION-REFERENCE IMAGES"}) + content.extend(references) + if historical: + content.append({"type": "text", "text": "TRAJECTORY AUDIT AND FINAL-CONTEXT IMAGES"}) + content.extend(_deduplicate_images(historical)) + if final: + content.append({"type": "text", "text": "FINAL/CURRENT OBSERVATIONS"}) + content.extend(_deduplicate_images(final)) + return content + + async def _judge_trace(self, trace: EpisodeTrace, complete: RoleCompleter) -> _JudgedTrace: + extra = _metadata(trace) + final_answer_raw = extra.get("final_answer") + final_answer = cast(dict[str, Any], final_answer_raw) if isinstance(final_answer_raw, dict) else {} + final_context = trace.rounds[-1].context_payload if trace.rounds else None + final_context_text, _ = _split_images(final_context) + environment_diagnostics = { + "environment_invalid": trace.metadata.environment_invalid, + "termination_reason": trace.metadata.termination_reason, + "runtime_errors": [*trace.metadata.runtime_errors, *trace.adapter_errors], + } + user_text = self.judger_user.format( + task_family=_family(trace), + scene=str(extra.get("scene", "")), + question_id=str(extra.get("question_id", "")), + question=str(extra.get("question", "")), + options=_json_text(extra.get("options")), + task_contract=str(extra.get("task_contract", "")), + final_answer=str(final_answer.get("answer", "not sure")), + final_confidence=final_answer.get("confidence", 0.0), + termination=trace.metadata.termination_reason, + ground_truth=_json_text(extra.get("ground_truth")), + official_correct=str(_official_correct(trace)).lower(), + environment_diagnostics=_json_text(environment_diagnostics), + planner_prompt=str(extra.get("planner_skill", "")), + final_context=_json_text(final_context_text), + trajectory=_trajectory(trace), + ) + content: list[dict[str, Any]] = [{"type": "text", "text": user_text}] + content.extend(self._judger_images(trace)) + raw = await complete( + RoleRequest( + system_prompt=self.judger_system, + user_content=content, + temperature=1.0, + response_format="json_object", + ) + ) + payload = parse_json_object(raw) + authoritative = _official_correct(trace) + run_valid = not trace.metadata.environment_invalid + payload["official_correct"] = authoritative + payload["run_valid"] = run_valid + score = payload.get("evidence_progress_score") + if not run_valid: + payload["evidence_progress_score"] = None + elif authoritative: + payload["evidence_progress_score"] = 1.0 + elif not isinstance(score, (int, float)) or float(score) not in {0.0, 0.25, 0.5, 0.75}: + raise ValueError(f"ESI-Bench Judger returned invalid evidence_progress_score: {score!r}") + else: + payload["evidence_progress_score"] = float(score) + for key in ( + "evidence_analysis", + "exploration_analysis", + "answer_analysis", + "context_analysis", + "improvement_signal", + ): + if not isinstance(payload.get(key), str): + raise ValueError(f"ESI-Bench Judger omitted string field {key!r}.") + failure = str(payload.get("primary_failure", "other")) + payload["primary_failure"] = failure if failure in _FAILURE_LABELS else "other" + target = str(payload.get("optimizer_target", "none")) + payload["optimizer_target"] = target if target in {"prompt", "context", "environment", "none"} else "none" + if not run_valid: + payload["primary_failure"] = "environment_failure" + payload["optimizer_target"] = "environment" + return _JudgedTrace(trace=trace, judgement=payload) + + @staticmethod + def _category_summary(evaluation: CandidateEvaluation) -> str: + totals: dict[str, int] = defaultdict(int) + correct: dict[str, int] = defaultdict(int) + invalid: dict[str, int] = defaultdict(int) + for trace in evaluation.traces: + family = _family(trace) + if trace.metadata.environment_invalid: + invalid[family] += 1 + else: + totals[family] += 1 + correct[family] += int(_official_correct(trace)) + families = sorted(set(totals) | set(invalid)) + return "\n".join( + f"{family}: {correct[family]}/{totals[family]} correct; {invalid[family]} environment-invalid" + for family in families + ) + + @staticmethod + def _execution_statistics(evaluation: CandidateEvaluation, judged: Sequence[_JudgedTrace]) -> str: + traces = [item.trace for item in judged] + steps = [len(trace.rounds) for trace in traces] + actions = Counter(record.command for trace in traces for record in trace.rounds) + terminations = Counter(trace.metadata.termination_reason for trace in evaluation.traces) + invalid_actions = sum( + not bool(record.action_result.get("action_valid", True)) for trace in traces for record in trace.rounds + ) + failures = Counter(str(item.judgement.get("primary_failure", "other")) for item in judged) + payload = { + "mean_steps": (sum(steps) / len(steps)) if steps else 0.0, + "min_steps": min(steps, default=0), + "max_steps": max(steps, default=0), + "invalid_model_actions": invalid_actions, + "termination_frequencies": dict(terminations), + "action_frequencies": dict(actions.most_common(20)), + "judged_failure_frequencies": dict(failures), + } + return _json_text(payload) + + async def build_textual_gradient( + self, + evaluation: CandidateEvaluation, + complete: RoleCompleter, + ) -> _ESIBenchGradient: + raw = await asyncio.gather( + *(self._judge_trace(trace, complete) for trace in evaluation.traces), + return_exceptions=True, + ) + judged: list[_JudgedTrace] = [] + for trace, value in zip(evaluation.traces, raw): + if isinstance(value, BaseException): + run_valid = not trace.metadata.environment_invalid + judged.append( + _JudgedTrace( + trace=trace, + judgement={ + "official_correct": _official_correct(trace), + "run_valid": run_valid, + "evidence_progress_score": ( + (1.0 if _official_correct(trace) else 0.0) if run_valid else None + ), + "evidence_analysis": "Judger output unavailable.", + "exploration_analysis": "Judger output unavailable.", + "answer_analysis": "Judger output unavailable.", + "context_analysis": "Judger output unavailable.", + "primary_failure": "other" if run_valid else "environment_failure", + "optimizer_target": "none" if run_valid else "environment", + "improvement_signal": f"Diagnostic failure: {value}", + }, + ) + ) + else: + judged.append(value) + + valid = [trace for trace in evaluation.traces if not trace.metadata.environment_invalid] + valid_judged = [item for item in judged if not item.trace.metadata.environment_invalid] + n_correct = sum(_official_correct(trace) for trace in valid) + n_scored = len(valid) + n_invalid = len(evaluation.traces) - n_scored + category_summary = self._category_summary(evaluation) + execution_statistics = self._execution_statistics(evaluation, valid_judged) + judgements_text = "\n\n".join(_json_text(item.judgement) for item in judged) + summary_user = self.summarizer_user.format( + candidate_name=evaluation.candidate_version, + n_correct=n_correct, + n_scored=n_scored, + score=f"{n_correct / max(1, n_scored):.6f}", + n_valid=n_scored, + n_environment_invalid=n_invalid, + category_summary=category_summary, + execution_statistics=execution_statistics, + judgements=judgements_text, + ) + try: + batch_summary = await complete( + RoleRequest( + system_prompt=self.summarizer_system, + user_content=summary_user, + temperature=1.0, + response_format="text", + ) + ) + except (RuntimeError, ValueError) as exc: + batch_summary = ( + f"The candidate scored {n_correct}/{n_scored}; {n_invalid} trajectories were " + f"environment-invalid. Batch summarization failed: {exc}." + ) + return _ESIBenchGradient( + evaluation=evaluation, + judged=tuple(judged), + batch_summary=batch_summary, + n_correct=n_correct, + n_scored=n_scored, + n_invalid=n_invalid, + category_summary=category_summary, + execution_statistics=execution_statistics, + ) + + @staticmethod + def _representative(gradient: _ESIBenchGradient, limit: int = 6) -> list[_JudgedTrace]: + valid = [item for item in gradient.judged if not item.trace.metadata.environment_invalid] + selected: list[_JudgedTrace] = [] + seen_failures: set[str] = set() + for item in valid: + failure = str(item.judgement.get("primary_failure", "other")) + if not _official_correct(item.trace) and failure not in seen_failures: + selected.append(item) + seen_failures.add(failure) + if len(selected) >= limit: + return selected + for item in valid: + if item not in selected: + selected.append(item) + if len(selected) >= limit: + break + return selected + + @classmethod + def _task_interfaces(cls, gradient: _ESIBenchGradient) -> str: + values: list[dict[str, Any]] = [] + for item in cls._representative(gradient): + extra = _metadata(item.trace) + values.append( + { + "task_family": _family(item.trace), + "task_subfamily": extra.get("task_subfamily"), + "question": extra.get("question"), + "answer_options": extra.get("options"), + "reference_image_count": len(_reference_images(item.trace)), + "official_task_contract": extra.get("task_contract"), + } + ) + return _json_text(values) + + @classmethod + def _reports(cls, gradient: _ESIBenchGradient) -> str: + return "\n\n".join(_json_text(item.judgement) for item in cls._representative(gradient)) + + @classmethod + def _trace_excerpts(cls, gradient: _ESIBenchGradient) -> str: + excerpts: list[dict[str, Any]] = [] + for item in cls._representative(gradient): + indices = _selected_round_indices(item.trace.rounds, limit=5) + excerpts.append( + { + "task_family": _family(item.trace), + "official_correct": _official_correct(item.trace), + "termination": item.trace.metadata.termination_reason, + "steps": [ + { + "step": index + 1, + "reason_action": item.trace.rounds[index].planner_response, + "parsed_action": item.trace.rounds[index].command, + "action_result": item.trace.rounds[index].action_result, + } + for index in indices + ], + } + ) + return _json_text(excerpts) + + @classmethod + def _context_payloads(cls, gradient: _ESIBenchGradient) -> tuple[str, list[dict[str, Any]]]: + excerpts: list[dict[str, Any]] = [] + images: list[dict[str, Any]] = [] + for item in cls._representative(gradient, limit=4): + indices = _selected_round_indices(item.trace.rounds, limit=3) + steps: list[dict[str, Any]] = [] + for index in indices: + record = item.trace.rounds[index] + payload_text, payload_images = _split_images(record.context_payload) + steps.append( + { + "step": index + 1, + "observable_harness_input": record.harness_input, + "actual_ordered_context_payload": payload_text, + "payload_image_count": len(payload_images), + "planner_action": record.command, + } + ) + images.extend(payload_images) + excerpts.append( + { + "task_family": _family(item.trace), + "official_correct": _official_correct(item.trace), + "selected_steps": steps, + } + ) + return _json_text(excerpts), _deduplicate_images(images) + + @staticmethod + def _history(context: OptimizerRequestContext) -> str: + events = [event.model_dump(mode="json") for event in context.optimization_history[-30:]] + return _json_text(events) + + @staticmethod + def _artifact_feedback(context: OptimizerRequestContext) -> str: + recent_errors = [ + event.validation_error for event in context.optimization_history[-30:] if event.validation_error + ] + base = "The current artifact passed static validation and all configured smoke probes." + if recent_errors: + base += "\nRecent rejected-candidate feedback:\n- " + "\n- ".join(recent_errors[-8:]) + if context.validation_feedback: + base += context.validation_feedback + return base + + @classmethod + def _optimizer_images(cls, gradient: _ESIBenchGradient, stage: str) -> list[dict[str, Any]]: + images: list[dict[str, Any]] + if stage == "harness": + _, images = cls._context_payloads(gradient) + label = "ACTUAL IMAGES FROM THE REPRESENTATIVE CONTEXT PAYLOADS" + else: + images = [] + for item in cls._representative(gradient, limit=4): + images.extend(_reference_images(item.trace)) + if item.trace.rounds: + images.extend(_image_parts(item.trace.rounds[0].observation_before)) + images.extend(_image_parts(item.trace.rounds[-1].observation_after)) + label = "REPRESENTATIVE OFFICIAL VISUAL EVIDENCE" + unique = _deduplicate_images(images) + return [{"type": "text", "text": label}, *unique] if unique else [] + + def build_optimizer_request(self, context: OptimizerRequestContext) -> RoleRequest: + gradient = cast(_ESIBenchGradient, context.textual_gradient) + common = { + "candidate_name": gradient.evaluation.candidate_version, + "n_correct": gradient.n_correct, + "n_scored": gradient.n_scored, + "score": f"{gradient.n_correct / max(1, gradient.n_scored):.6f}", + "n_valid": gradient.n_scored, + "n_environment_invalid": gradient.n_invalid, + "batch_summary": gradient.batch_summary, + "representative_task_interfaces": self._task_interfaces(gradient), + "representative_judger_reports": self._reports(gradient), + "representative_trace_excerpts": self._trace_excerpts(gradient), + "optimization_history": self._history(context), + } + if context.stage == "skill": + user_text = self.skill_optimizer_user.format( + current_skill=context.parent.skill.template, + frozen_harness=context.parent.harness.template, + **common, + ) + system_prompt = self.skill_optimizer_system + else: + context_payloads, _ = self._context_payloads(gradient) + user_text = self.harness_optimizer_user.format( + frozen_planner_skill=context.parent.skill.template, + current_context_code=context.parent.harness.template, + representative_context_payloads=context_payloads, + artifact_validation_feedback=self._artifact_feedback(context), + **common, + ) + user_text += ( + "\n\nAGENT LIGHTNING HARNESS VALIDATION CONTRACT\n" + "===========================================\n" + + context.harness_contract + + f"\nFunction: {context.harness_function_name}" + + f"\nSmoke arguments: {list(context.harness_smoke_args)!r}" + ) + system_prompt = self.harness_optimizer_system + "\n\n" + _AGL_HARNESS_ADAPTER + if context.validation_feedback and context.stage == "skill": + user_text += context.validation_feedback + user_text += f"\n\nProposal round: {context.round_index}; branch: {context.branch_index}." + content: list[dict[str, Any]] = [{"type": "text", "text": user_text}] + content.extend(self._optimizer_images(gradient, context.stage)) + return RoleRequest( + system_prompt=system_prompt, + user_content=content, + temperature=1.0, + response_format="json_object", + ) + + def parse_optimizer_response(self, stage: OptimizationStage, response: str) -> ArtifactProposal: + del stage + payload = parse_json_object(response) + rationale = payload.get("rationale") + artifact = payload.get("new_artifact") + if not isinstance(rationale, str) or not isinstance(artifact, str) or not artifact.strip(): + raise ValueError("ESI-Bench optimizer must return string fields rationale and new_artifact.") + return ArtifactProposal(rationale=rationale, artifact=artifact.strip()) + + +__all__ = ["ESIBenchRoleProtocol"] diff --git a/contrib/recipes/shaper/esi_bench/splits/README.md b/contrib/recipes/shaper/esi_bench/splits/README.md new file mode 100644 index 000000000..e6de9a999 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/splits/README.md @@ -0,0 +1,14 @@ +# ESI-Bench Recipe Splits + +`recipe_train10.txt` and `recipe_validation10.txt` are deterministic, disjoint +development manifests supplied so the contrib recipe runs without untracked +files. Each contains one question from every official top-level category. Both +are disjoint from `reported_eval231.txt`, the evaluation manifest associated +with the author-reported result snapshot. + +These files are not claimed to be the exact 10/10 manifests used to produce the +paper tables because those original manifests were not recoverable. Set +`ESI_TRAIN_SPLIT` and `ESI_VALIDATION_SPLIT` to explicit replacements when +reproducing another run. `recipe_metadata.json` records the selected task +families and pinned upstream commit without copying answers or simulator +metadata. diff --git a/contrib/recipes/shaper/esi_bench/splits/recipe_metadata.json b/contrib/recipes/shaper/esi_bench/splits/recipe_metadata.json new file mode 100644 index 000000000..a9d5c4b54 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/splits/recipe_metadata.json @@ -0,0 +1,29 @@ +{ + "upstream_repository": "https://github.com/ESI-Bench/ESI-Bench", + "upstream_commit": "3c1756396f32b1a90c1f72356a7fde45f418e179", + "status": "deterministic recipe split; not claimed as the paper's original 10/10 manifest", + "train": [ + {"id": "0028", "big_task": "Action Sequencing", "small_task": "Action Order Inference", "runner_task": "action"}, + {"id": "0197", "big_task": "Cognitive Mapping", "small_task": "Long-Term Navigation", "runner_task": "cognitivemap"}, + {"id": "0460", "big_task": "Enumerative Perception", "small_task": "Illumination Variability", "runner_task": "counting"}, + {"id": "0636", "big_task": "Metric Comparison", "small_task": "Dimensional Size", "runner_task": "size"}, + {"id": "1025", "big_task": "Perceptual Grounding", "small_task": "Material Transparency", "runner_task": "transparent"}, + {"id": "1826", "big_task": "Physical Dynamics", "small_task": "Stacking & Stability", "runner_task": "stacking"}, + {"id": "1847", "big_task": "Physical Structure", "small_task": "Deformable", "runner_task": "deformable"}, + {"id": "2905", "big_task": "Spatial Relations", "small_task": "Physical Contact", "runner_task": "touching"}, + {"id": "3071", "big_task": "Specular Reflection", "small_task": "Reflection Authoring", "runner_task": "mirror"}, + {"id": "3287", "big_task": "Temporal Understanding", "small_task": "Agent Observation", "runner_task": "multiagent"} + ], + "validation": [ + {"id": "0001", "big_task": "Action Sequencing", "small_task": "Action Order Inference", "runner_task": "action"}, + {"id": "0078", "big_task": "Cognitive Mapping", "small_task": "Connectivity", "runner_task": "cognitivemap"}, + {"id": "0338", "big_task": "Enumerative Perception", "small_task": "Category Ambiguity", "runner_task": "counting"}, + {"id": "0608", "big_task": "Metric Comparison", "small_task": "Dimensional Size", "runner_task": "size"}, + {"id": "1171", "big_task": "Perceptual Grounding", "small_task": "Partial Occlusion", "runner_task": "occlusion"}, + {"id": "1731", "big_task": "Physical Dynamics", "small_task": "Inclined Plane", "runner_task": "slope"}, + {"id": "2061", "big_task": "Physical Structure", "small_task": "Rigid Containment", "runner_task": "storage"}, + {"id": "2394", "big_task": "Spatial Relations", "small_task": "Geometric Configuration", "runner_task": "triangle"}, + {"id": "3166", "big_task": "Specular Reflection", "small_task": "Spatial Relations", "runner_task": "mirror"}, + {"id": "3422", "big_task": "Temporal Understanding", "small_task": "Unobserved Change", "runner_task": "unobserved_changes"} + ] +} diff --git a/contrib/recipes/shaper/esi_bench/splits/recipe_train10.txt b/contrib/recipes/shaper/esi_bench/splits/recipe_train10.txt new file mode 100644 index 000000000..7d2eb6d4f --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/splits/recipe_train10.txt @@ -0,0 +1,10 @@ +0028 +0197 +0460 +0636 +1025 +1826 +1847 +2905 +3071 +3287 diff --git a/contrib/recipes/shaper/esi_bench/splits/recipe_validation10.txt b/contrib/recipes/shaper/esi_bench/splits/recipe_validation10.txt new file mode 100644 index 000000000..0905b8398 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/splits/recipe_validation10.txt @@ -0,0 +1,10 @@ +0001 +0078 +0338 +0608 +1171 +1731 +2061 +2394 +3166 +3422 diff --git a/contrib/recipes/shaper/esi_bench/splits/reported_eval231.txt b/contrib/recipes/shaper/esi_bench/splits/reported_eval231.txt new file mode 100644 index 000000000..9aa3a98a1 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/splits/reported_eval231.txt @@ -0,0 +1,231 @@ +0016 +0017 +0041 +0070 +0074 +0098 +0103 +0111 +0160 +0201 +0207 +0208 +0210 +0212 +0217 +0228 +0229 +0236 +0239 +0241 +0261 +0262 +0348 +0357 +0378 +0388 +0391 +0393 +0399 +0400 +0409 +0413 +0417 +0425 +0510 +0519 +0531 +0548 +0559 +0583 +0648 +0658 +0661 +0671 +0675 +0701 +0707 +0715 +0717 +0750 +0751 +0759 +0765 +0784 +0802 +0806 +0811 +0822 +0882 +0888 +0897 +0938 +0940 +0962 +0965 +0977 +0979 +1012 +1016 +1018 +1026 +1033 +1036 +1046 +1051 +1058 +1060 +1061 +1066 +1068 +1073 +1098 +1126 +1137 +1138 +1153 +1164 +1172 +1174 +1189 +1194 +1199 +1210 +1213 +1229 +1231 +1271 +1287 +1339 +1351 +1357 +1360 +1361 +1373 +1445 +1451 +1555 +1556 +1580 +1592 +1695 +1700 +1716 +1720 +1728 +1750 +1756 +1757 +1766 +1782 +1852 +1875 +1883 +1897 +2065 +2067 +2071 +2080 +2082 +2085 +2087 +2097 +2101 +2103 +2106 +2109 +2110 +2124 +2125 +2126 +2129 +2224 +2227 +2245 +2272 +2304 +2305 +2313 +2318 +2333 +2361 +2368 +2407 +2413 +2419 +2439 +2476 +2480 +2490 +2520 +2536 +2625 +2630 +2650 +2679 +2684 +2693 +2694 +2714 +2728 +2730 +2734 +2761 +2762 +2763 +2764 +2766 +2775 +2787 +2789 +2791 +2808 +2828 +2830 +2834 +2840 +2865 +2868 +2883 +2895 +2902 +2904 +2930 +2937 +2939 +2948 +2949 +2951 +2954 +2973 +2985 +2989 +2990 +2993 +3000 +3003 +3006 +3008 +3027 +3039 +3043 +3114 +3256 +3262 +3266 +3270 +3280 +3293 +3294 +3295 +3303 +3317 +3346 +3350 +3351 +3376 +3399 +3429 +3439 +3473 +3478 diff --git a/contrib/recipes/shaper/esi_bench/train.py b/contrib/recipes/shaper/esi_bench/train.py new file mode 100644 index 000000000..333469a68 --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/train.py @@ -0,0 +1,68 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Preflight and run SHAPER training on the included ESI-Bench adapter.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path +from typing import Sequence + +from ..cli import cli_arguments, print_preflight_errors, requests_help, required_environment +from ..reproduce import main as reproduce_main +from .check_env import absolute_executable, check_environment + +FACTORY = "contrib.recipes.shaper.esi_bench.factory:build_bundle" +RECIPE_DIR = Path(__file__).parent + + +def preflight_environment() -> list[str]: + root = Path(required_environment("ESI_BENCH_ROOT")).expanduser().resolve() + behavior_root = Path(required_environment("ESI_BEHAVIOR_ROOT")).expanduser().resolve() + raw_data_root = os.environ.get("ESI_OMNIGIBSON_DATA_ROOT") or os.environ.get("OMNIGIBSON_DATA_PATH") + if not raw_data_root: + raise ValueError("Set ESI_OMNIGIBSON_DATA_ROOT or OMNIGIBSON_DATA_PATH before running this command.") + omnigibson_data_root = Path(raw_data_root).expanduser().resolve() + questions = ( + Path(os.environ.get("ESI_QUESTIONS_JSONL", root / "hf_dataset" / "data" / "questions.jsonl")) + .expanduser() + .resolve() + ) + train_split = ( + Path(os.environ.get("ESI_TRAIN_SPLIT", RECIPE_DIR / "splits" / "recipe_train10.txt")).expanduser().resolve() + ) + validation_split = ( + Path(os.environ.get("ESI_VALIDATION_SPLIT", RECIPE_DIR / "splits" / "recipe_validation10.txt")) + .expanduser() + .resolve() + ) + make_maps_path = Path(required_environment("ESI_MAKE_MAPS_PATH")).expanduser().resolve() + worker_python = absolute_executable(Path(os.environ.get("ESI_WORKER_PYTHON", sys.executable))) + return check_environment( + root=root, + behavior_root=behavior_root, + omnigibson_data_root=omnigibson_data_root, + questions_jsonl=questions, + train_split=train_split, + validation_split=validation_split, + make_maps_path=make_maps_path, + worker_python=worker_python, + planner_endpoint=required_environment("SHAPER_PLANNER_ENDPOINT"), + require_planner=True, + ) + + +def main(argv: Sequence[str] | None = None) -> int: + arguments = cli_arguments(argv) + if requests_help(arguments): + reproduce_main(["--factory", FACTORY, *arguments]) + return 0 + if print_preflight_errors(preflight_environment()): + return 2 + reproduce_main(["--factory", FACTORY, *arguments]) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/contrib/recipes/shaper/esi_bench/worker.py b/contrib/recipes/shaper/esi_bench/worker.py new file mode 100644 index 000000000..4799c950d --- /dev/null +++ b/contrib/recipes/shaper/esi_bench/worker.py @@ -0,0 +1,1052 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Fresh-process bridge to ESI-Bench's official active-exploration pipeline.""" + +from __future__ import annotations + +import argparse +import base64 +import importlib +import json +import logging +import os +import re +import sys +from pathlib import Path +from typing import Any, Callable, Mapping, Sequence, cast + +from openai import OpenAI + +from ..common import ( + first_json_object, + image_part, + normalize_content, + path_data_url, + sanitized_action_result, + strip_thinking, + text_part, +) +from ..harness_bridge import HarnessBridgeClient, HarnessBridgeError +from .contracts import ( + check_behavior_source, + check_omnigibson_install, + check_upstream_source, + validate_skill, +) +from .dataset import load_question_row, resolve_canonical_question + +logger = logging.getLogger(__name__) + +GENERIC_JSON_INSTRUCTION = "Return exactly one valid JSON object and nothing else." + +_deferred_simulator_shutdown: Callable[[], Any] | None = None + + +class PlannerRequestError(RuntimeError): + """The frozen planner endpoint failed before an action could be produced.""" + + +class OfficialRunnerError(RuntimeError): + """The official simulator or task runner failed independently of artifacts.""" + + +def _render_skill(skill: str, task_prompt: str) -> str: + """Wrap one official task prompt with the validated reusable skill.""" + + if skill.count("{task_prompt}") != 1: + raise ValueError("ESI-Bench skill must contain {task_prompt} exactly once.") + return skill.replace("{task_prompt}", task_prompt, 1) + + +def _jpeg_data_url(image: Any, *, max_side: int = 768, quality: int = 90) -> str: + """Encode a bounded deterministic BGR array for harness transport.""" + + cv2 = cast(Any, importlib.import_module("cv2")) + numpy = cast(Any, importlib.import_module("numpy")) + encoded_image = numpy.asarray(image) + height, width = encoded_image.shape[:2] + if max(width, height) > max_side: + scale = max_side / float(max(width, height)) + target = (max(1, round(width * scale)), max(1, round(height * scale))) + encoded_image = cv2.resize(encoded_image, target, interpolation=cv2.INTER_AREA) + ok, encoded = cv2.imencode( + ".jpg", + numpy.ascontiguousarray(encoded_image), + [int(cv2.IMWRITE_JPEG_QUALITY), quality], + ) + if not ok: + raise ValueError("OpenCV could not encode the official ESI-Bench RGB.") + return "data:image/jpeg;base64," + base64.b64encode(encoded.tobytes()).decode("ascii") + + +def _pixel_quality(image: Any) -> dict[str, float]: + """Score one crop using rendered pixels only. + + The three terms mirror the evidence-quality signals described by SHAPER: + grayscale contrast, finite-difference edge density, and Laplacian + sharpness. They carry no semantic label or simulator metadata. + """ + + numpy = cast(Any, importlib.import_module("numpy")) + pixels = numpy.asarray(image) + if pixels.size == 0: + return {"contrast": 0.0, "edge_density": 0.0, "laplacian_sharpness": 0.0, "score": 0.0} + if pixels.ndim == 3 and pixels.shape[2] >= 3: + gray = 0.114 * pixels[..., 0] + 0.587 * pixels[..., 1] + 0.299 * pixels[..., 2] + elif pixels.ndim == 3: + gray = pixels.mean(axis=2) + else: + gray = pixels + gray = numpy.asarray(gray, dtype=numpy.float32) + row_stride = max(1, (int(gray.shape[0]) + 71) // 72) + column_stride = max(1, (int(gray.shape[1]) + 127) // 128) + gray = gray[::row_stride, ::column_stride][:72, :128] + if gray.shape[0] < 3 or gray.shape[1] < 3: + return {"contrast": 0.0, "edge_density": 0.0, "laplacian_sharpness": 0.0, "score": 0.0} + contrast = min(1.5, float(gray.std()) / 55.0) + horizontal = numpy.abs(numpy.diff(gray, axis=1)) + vertical = numpy.abs(numpy.diff(gray, axis=0)) + edge_density = 0.5 * (float((horizontal > 16.0).mean()) + float((vertical > 16.0).mean())) + center = gray[1:-1, 1:-1] + laplacian = gray[:-2, 1:-1] + gray[2:, 1:-1] + gray[1:-1, :-2] + gray[1:-1, 2:] - 4.0 * center + sharpness = min(1.5, float(laplacian.var()) / 900.0) + score = contrast + 5.0 * edge_density + 0.35 * sharpness + return { + "contrast": round(contrast, 6), + "edge_density": round(edge_density, 6), + "laplacian_sharpness": round(sharpness, 6), + "score": round(score, 6), + } + + +def _focus_crop_candidates(image: Any) -> list[dict[str, Any]]: + """Select one horizontal band and one overlapping tile deterministically.""" + + height, width = image.shape[:2] + candidates: list[tuple[str, str, Any, dict[str, float]]] = [] + band_height = min(height, max(1, round(height * 0.28))) + for index, fraction in enumerate((0.0, 0.12, 0.24, 0.36, 0.48, 0.60, 0.72), start=1): + top = min(height - band_height, round(height * fraction)) + crop = image[top : top + band_height, 0:width] + candidates.append(("horizontal_band", f"horizontal_band_{index}", crop, _pixel_quality(crop))) + + tile_width = min(width, max(1, round(width * 0.62))) + tile_height = min(height, max(1, round(height * 0.55))) + tile_index = 0 + for y_fraction in (0.0, 0.225, 0.45): + for x_fraction in (0.0, 0.19, 0.38): + tile_index += 1 + left = min(width - tile_width, round(width * x_fraction)) + top = min(height - tile_height, round(height * y_fraction)) + crop = image[top : top + tile_height, left : left + tile_width] + candidates.append(("overlapping_tile", f"overlapping_tile_{tile_index}", crop, _pixel_quality(crop))) + + selected: list[dict[str, Any]] = [] + for family in ("horizontal_band", "overlapping_tile"): + family_candidates = [candidate for candidate in candidates if candidate[0] == family] + _, region, crop, quality = max( + family_candidates, + key=lambda candidate: (float(candidate[3]["score"]), candidate[1]), + ) + selected.append( + { + "source": "deterministic_pixel_crop", + "region": region, + "quality": quality, + "image": image_part(_jpeg_data_url(crop, max_side=512)), + } + ) + return selected + + +def _visual_signature(image: Any) -> list[int]: + """Return a tiny pixel-only signature for bounded keyframe diversity.""" + + numpy = cast(Any, importlib.import_module("numpy")) + pixels = numpy.asarray(image) + if pixels.ndim == 3 and pixels.shape[2] >= 3: + gray = 0.114 * pixels[..., 0] + 0.587 * pixels[..., 1] + 0.299 * pixels[..., 2] + elif pixels.ndim == 3: + gray = pixels.mean(axis=2) + else: + gray = pixels + gray = numpy.asarray(gray, dtype=numpy.float32) + if gray.size == 0: + return [0] * 64 + rows = numpy.rint(numpy.linspace(0, gray.shape[0] - 1, 8)).astype(int) + columns = numpy.rint(numpy.linspace(0, gray.shape[1] - 1, 8)).astype(int) + reduced = numpy.clip(numpy.rint(gray[numpy.ix_(rows, columns)]), 0, 255).astype(int) + return [int(value) for value in reduced.reshape(-1).tolist()] + + +def _grid_overlay(image: Any) -> dict[str, Any]: + """Overlay a deterministic GRID1000 coordinate system on official pixels.""" + + cv2 = cast(Any, importlib.import_module("cv2")) + numpy = cast(Any, importlib.import_module("numpy")) + grid = numpy.asarray(image).copy() + height, width = grid.shape[:2] + line_color = (214, 160, 64) + text_color = (255, 255, 255) + shadow_color = (24, 24, 24) + thickness = max(1, round(max(width, height) / 900)) + font_scale = max(0.35, min(width, height) / 1300.0) + for value in range(0, 1001, 100): + x = min(width - 1, round((width - 1) * value / 1000.0)) + y = min(height - 1, round((height - 1) * value / 1000.0)) + major = value % 250 == 0 + cv2.line(grid, (x, 0), (x, height - 1), line_color, thickness + int(major)) + cv2.line(grid, (0, y), (width - 1, y), line_color, thickness + int(major)) + if major: + x_origin = min(max(2, x + 3), max(2, width - 46)) + y_origin = min(max(14, y + 14), max(14, height - 4)) + label = str(value) + cv2.putText( + grid, + label, + (x_origin + 1, 14), + cv2.FONT_HERSHEY_SIMPLEX, + font_scale, + shadow_color, + thickness + 2, + cv2.LINE_AA, + ) + cv2.putText( + grid, + label, + (x_origin, 13), + cv2.FONT_HERSHEY_SIMPLEX, + font_scale, + text_color, + thickness, + cv2.LINE_AA, + ) + cv2.putText( + grid, + label, + (3, y_origin + 1), + cv2.FONT_HERSHEY_SIMPLEX, + font_scale, + shadow_color, + thickness + 2, + cv2.LINE_AA, + ) + cv2.putText( + grid, + label, + (2, y_origin), + cv2.FONT_HERSHEY_SIMPLEX, + font_scale, + text_color, + thickness, + cv2.LINE_AA, + ) + return { + "source": "deterministic_pixel_overlay", + "coordinate_system": "GRID1000: x=0..1000 left-to-right, y=0..1000 top-to-bottom", + "image": image_part(_jpeg_data_url(grid)), + } + + +def _needs_geometry_grid(task_prompt: str) -> bool: + """Select pixel overlays from the official visible task contract only.""" + + text = task_prompt.lower() + return bool( + any(term in text for term in ("triangle", "equilateral", "isosceles", "collinear")) + or re.search(r"\b(?:in a|straight) line\b", text) + ) + + +class _ObservableImageCache: + """Create full-frame and pixel-ranked crop payloads from official RGBs.""" + + def __init__(self) -> None: + self._cache: dict[str, dict[str, Any]] = {} + self._grid_cache: dict[str, dict[str, Any]] = {} + + def observation(self, path: Path, *, include_grid: bool = False) -> dict[str, Any]: + key = str(path.resolve()) + cached = self._cache.get(key) + if cached is None: + if not path.is_file(): + raise FileNotFoundError(f"Official ESI-Bench RGB does not exist: {path}") + try: + cv2 = cast(Any, importlib.import_module("cv2")) + image = cv2.imread(str(path), cv2.IMREAD_COLOR) + if image is None: + raise ValueError("OpenCV could not decode the image.") + height, width = image.shape[:2] + value: dict[str, Any] = { + "source": "official_rgb", + "pixel_size": {"width": int(width), "height": int(height)}, + "pixel_quality": _pixel_quality(image), + "visual_signature": { + "kind": "8x8_grayscale", + "values": _visual_signature(image), + }, + "full_frame": image_part(_jpeg_data_url(image)), + "focus_crops": _focus_crop_candidates(image), + } + except Exception as exc: + logger.warning("Falling back to the original RGB bytes for %s: %s", path, exc) + value = { + "source": "official_rgb", + "pixel_size": None, + "pixel_quality": None, + "visual_signature": None, + "full_frame": image_part(path_data_url(path)), + "focus_crops": [], + } + self._cache[key] = value + cached = value + output = cast(dict[str, Any], json.loads(json.dumps(cached))) + if include_grid: + grid = self._grid_cache.get(key) + if grid is None: + try: + cv2 = cast(Any, importlib.import_module("cv2")) + image = cv2.imread(str(path), cv2.IMREAD_COLOR) + if image is not None: + grid = _grid_overlay(image) + self._grid_cache[key] = grid + except Exception as exc: + logger.warning("Could not construct GRID1000 overlay for %s: %s", path, exc) + if grid is not None: + output["grid_overlay"] = cast( + dict[str, Any], + json.loads(json.dumps(grid)), + ) + return output + + +def _extra_paths(item: Mapping[str, Any]) -> list[Path]: + output: list[Path] = [] + raw = item.get("extra_image_paths") + if isinstance(raw, list): + for value in cast(list[Any], raw): + if isinstance(value, (str, Path)): + output.append(Path(value)) + return output + + +def _without_inline_pixels(value: Any) -> Any: + """Keep harness-input provenance while avoiding duplicate base64 payloads.""" + + if isinstance(value, list): + return [_without_inline_pixels(item) for item in cast(list[Any], value)] + if not isinstance(value, dict): + return value + mapping = cast(dict[str, Any], value) + if mapping.get("type") == "image_url" and isinstance(mapping.get("image_url"), dict): + return {"type": "image_url", "image_url": {"url": ""}} + return {str(key): _without_inline_pixels(item) for key, item in mapping.items()} + + +def _decoded_json_field(row: Mapping[str, Any], key: str) -> Any: + value = row.get(key) + if not isinstance(value, str): + return value + try: + return json.loads(value) + except json.JSONDecodeError: + return value + + +class _HarnessCollector: + """Route every official planner context through one restricted harness.""" + + def __init__( + self, + skill: str, + build_context: Callable[[list[dict[str, Any]]], Any], + *, + max_steps: int, + ) -> None: + self.skill = skill + self.build_context = build_context + self.max_steps = max_steps + self.images = _ObservableImageCache() + self.snapshots: dict[int, dict[str, Any]] = {} + self.reference_images: list[dict[str, Any]] = [] + self._latest_past_records: list[dict[str, Any]] = [] + self._latest_history_index: int | None = None + + def _past_record( + self, + image_dir: Path, + item: Mapping[str, Any], + *, + include_grid: bool, + ) -> dict[str, Any]: + image_path = image_dir / str(item.get("image", "")) + extras = [ + self.images.observation(path, include_grid=include_grid) for path in _extra_paths(item) if path.is_file() + ] + action_result = sanitized_action_result(item.get("action_result")) + return { + "record_kind": "past", + "step": int(item.get("step", 0)), + "action": str(item.get("action", "")), + "answer": str(item.get("answer", "")), + "confidence": float(item.get("confidence", 0.0)), + "reasoning": str(item.get("reasoning", "")), + "action_result": action_result, + "action_result_text": json.dumps(action_result, ensure_ascii=True), + "observation": self.images.observation(image_path, include_grid=include_grid), + "extra_observations": extras, + } + + def __call__( + self, + image_path: Path, + history: list[dict[str, Any]], + prompt: str, + reference_image_paths: list[Path] | None = None, + reference_image_path: Path | None = None, + ) -> list[dict[str, Any]]: + references = list(reference_image_paths or []) + if reference_image_path is not None: + references.append(reference_image_path) + include_grid = _needs_geometry_grid(prompt) + records = [self._past_record(image_path.parent, item, include_grid=include_grid) for item in history] + reference_count = len(references) + current: dict[str, Any] = { + "record_kind": "current", + "step": len(history) + 1, + "max_steps": self.max_steps, + "remaining_steps": max(0, self.max_steps - len(history)), + "task_instruction": prompt, + "observation": self.images.observation(image_path, include_grid=include_grid), + "reference_observations": [ + { + "label": ( + "[QUESTION REFERENCE IMAGE - dataset render]" + if reference_count == 1 + else "[QUESTION REFERENCE IMAGE " + str(index) + "]" + ), + **self.images.observation(path, include_grid=include_grid), + } + for index, path in enumerate(references, start=1) + if path.is_file() + ], + } + records.append(current) + if not self.reference_images: + self.reference_images = [ + { + "label": str(reference.get("label", "QUESTION REFERENCE IMAGE")), + "image": cast(dict[str, Any], reference["full_frame"]), + } + for reference in cast(list[dict[str, Any]], current["reference_observations"]) + if isinstance(reference.get("full_frame"), dict) + ] + context = normalize_content(self.build_context(records)) + self._latest_past_records = records[:-1] + self._latest_history_index = len(history) + current_observation = cast(dict[str, Any], current["observation"]) + self.snapshots.setdefault( + len(history), + { + "task_instruction": prompt, + "harness_input": _without_inline_pixels(records), + "context_payload": context, + "observation_before": [current_observation["full_frame"]], + }, + ) + # Match the official collector ordering: visual/history context first, + # then the authoritative task prompt (wrapped by the selected skill). + return [*context, text_part(_render_skill(self.skill, prompt))] + + def route_auxiliary( + self, + contents: Sequence[Any], + task_prompt: str, + ) -> list[dict[str, Any]]: + """Route one audited task-specific call through the selected harness. + + The pinned inclined-plane hook supplies only ordered official RGB paths + and visible text. They remain in that order, but RGBs receive the same + pixel-only derivatives as primary observations so an evolved harness + can select evidence instead of inheriting an unchangeable side path. + """ + + include_grid = _needs_geometry_grid(task_prompt) + sequence: list[dict[str, Any]] = [] + for index, item in enumerate(contents): + if isinstance(item, Path): + sequence.append( + { + "content_kind": "observation", + "sequence_index": index, + "observation": self.images.observation(item, include_grid=include_grid), + } + ) + elif isinstance(item, dict): + mapping = cast(dict[str, Any], item) + if mapping.get("type") == "text" and isinstance(mapping.get("text"), str): + sequence.append( + { + "content_kind": "text", + "sequence_index": index, + "text": str(mapping["text"]), + } + ) + elif mapping.get("type") == "image_url" and isinstance(mapping.get("image_url"), dict): + sequence.append( + { + "content_kind": "content_part", + "sequence_index": index, + "part": mapping, + } + ) + else: + sequence.append( + { + "content_kind": "text", + "sequence_index": index, + "text": str(mapping), + } + ) + else: + sequence.append( + { + "content_kind": "text", + "sequence_index": index, + "text": str(item), + } + ) + + history_index = self._latest_history_index + step = (history_index + 1) if history_index is not None else 1 + current: dict[str, Any] = { + "record_kind": "current", + "call_kind": "auxiliary_post_action", + "step": step, + "max_steps": self.max_steps, + "remaining_steps": max(0, self.max_steps - (history_index or 0)), + "task_instruction": task_prompt, + "observable_sequence": sequence, + "reference_observations": [], + } + records = [*self._latest_past_records, current] + context = normalize_content(self.build_context(records)) + + if history_index is not None: + snapshot = self.snapshots.setdefault( + history_index, + { + "task_instruction": task_prompt, + "harness_input": _without_inline_pixels(records), + "context_payload": [], + "observation_before": [], + }, + ) + existing = snapshot.get("context_payload") + existing_contexts = cast(dict[str, Any], existing) if isinstance(existing, dict) else None + call_contexts: dict[str, Any] + if existing_contexts is not None and existing_contexts.get("call_contexts_version") == 1: + call_contexts = existing_contexts + else: + call_contexts = { + "call_contexts_version": 1, + "primary": existing, + "auxiliary_post_action": [], + } + snapshot["context_payload"] = call_contexts + auxiliary_value: Any = call_contexts.get("auxiliary_post_action") + auxiliary: list[Any] + if isinstance(auxiliary_value, list): + auxiliary = cast(list[Any], auxiliary_value) + else: + auxiliary = [] + call_contexts["auxiliary_post_action"] = auxiliary + auxiliary.append(context) + return context + + +def _partial_json(raw_text: str, fallback: Mapping[str, Any] | None) -> dict[str, Any]: + parsed = first_json_object(raw_text) + if parsed is not None: + return parsed + output = dict(fallback or {}) + for key in ("action", "answer", "reasoning"): + match = re.search(rf'"{key}"\s*:\s*"([^\"]*)', raw_text, flags=re.DOTALL) + if match: + output[key] = match.group(1).strip() + confidence = re.search(r'"confidence"\s*:\s*(-?\d+(?:\.\d+)?)', raw_text) + if confidence: + output["confidence"] = float(confidence.group(1)) + return output + + +class _OpenAIPlanner: + """OpenAI-compatible model client satisfying ESI-Bench's model protocol.""" + + def __init__( + self, + planner: Mapping[str, Any], + skill: str, + *, + auxiliary_context_builder: Callable[[Sequence[Any], str], list[dict[str, Any]]] | None = None, + ) -> None: + self.model = str(planner["model"]) + self.skill = skill + self.auxiliary_context_builder = auxiliary_context_builder + raw_sampling: object = planner.get("sampling_parameters") + self.sampling: dict[str, Any] = ( + dict(cast(Mapping[str, Any], raw_sampling)) if isinstance(raw_sampling, Mapping) else {} + ) + self.client = OpenAI( + api_key=str(planner.get("api_key") or "not-required"), + base_url=str(planner["endpoint"]), + timeout=float(self.sampling.get("timeout", 300.0)), + max_retries=int(self.sampling.get("max_retries", 2)), + ) + + @staticmethod + def _contents(contents: Sequence[Any]) -> list[dict[str, Any]]: + output: list[dict[str, Any]] = [] + for item in contents: + if isinstance(item, Path): + output.append(image_part(path_data_url(item))) + elif isinstance(item, dict): + mapping = cast(dict[str, Any], item) + if mapping.get("type") in {"text", "image_url"}: + output.append(mapping) + else: + output.append(text_part(mapping)) + else: + output.append(text_part(item)) + return output + + @staticmethod + def _unsupported_sampling_error(exc: BaseException) -> bool: + """Recognize a provider's explicit pre-generation sampling rejection.""" + + status = getattr(exc, "status_code", None) + text = str(exc).lower() + parameter = "temperature" in text or "top_p" in text + rejection = any(marker in text for marker in ("unsupported", "not supported", "unrecognized", "unknown")) + return status == 400 and parameter and rejection + + def generate_json( + self, + contents: list[Any], + system_instruction: str, + response_schema: dict[str, Any] | None = None, + max_output_tokens: int = 1024, + temperature: float = 1.0, + top_p: float = 0.95, + fallback: dict[str, Any] | None = None, + ) -> tuple[dict[str, Any], str, str | None]: + del response_schema + if system_instruction == GENERIC_JSON_INSTRUCTION: + system = system_instruction + routed_contents = contents + else: + if self.auxiliary_context_builder is None: + raise RuntimeError("Task-specific ESI-Bench planner calls require the selected context harness.") + system = _render_skill(self.skill, system_instruction) + routed_contents = self.auxiliary_context_builder(contents, system_instruction) + configured_max_tokens = int(self.sampling.get("max_completion_tokens", max_output_tokens)) + if configured_max_tokens < 1 or max_output_tokens < 1: + raise ValueError("Planner completion-token limits must be positive.") + request: dict[str, Any] = { + "model": self.model, + "messages": [ + {"role": "system", "content": system}, + {"role": "user", "content": self._contents(routed_contents)}, + ], + # The deployment-wide setting is a ceiling, not permission to + # override the official runner's smaller auxiliary/final limits. + "max_completion_tokens": min(configured_max_tokens, max_output_tokens), + "response_format": {"type": "json_object"}, + } + request["temperature"] = self.sampling.get("temperature", temperature) + request["top_p"] = self.sampling.get("top_p", top_p) + if "presence_penalty" in self.sampling: + request["presence_penalty"] = self.sampling["presence_penalty"] + extra_body: object = self.sampling.get("extra_body") + if isinstance(extra_body, dict): + request["extra_body"] = cast(dict[str, Any], extra_body) + try: + try: + response = cast(Any, self.client.chat.completions.create(**request)) + except Exception as exc: + if not self._unsupported_sampling_error(exc): + raise + request.pop("temperature", None) + request.pop("top_p", None) + response = cast(Any, self.client.chat.completions.create(**request)) + except Exception as exc: + raise PlannerRequestError(f"Frozen planner request failed: {exc}") from exc + choice = response.choices[0] + raw_text = strip_thinking(choice.message.content or "") + return ( + _partial_json(raw_text, fallback), + raw_text, + str(getattr(choice, "finish_reason", None)), + ) + + +def _load_official_pipeline(root: Path, behavior_root: Path) -> Any: + """Import the pinned official runner only inside the simulator process.""" + + source_errors = check_upstream_source(root) + if source_errors: + raise RuntimeError("Unsupported ESI-Bench checkout: " + "; ".join(source_errors)) + behavior_errors = [*check_behavior_source(behavior_root), *check_omnigibson_install(behavior_root)] + if behavior_errors: + raise RuntimeError("Unsupported BEHAVIOR/OmniGibson environment: " + "; ".join(behavior_errors)) + active_root = root / "src" / "active_explore" + pipeline_path = active_root / "pipeline.py" + if not pipeline_path.is_file(): + raise FileNotFoundError(f"Official ESI-Bench pipeline not found: {pipeline_path}") + sys.path.insert(0, str(active_root)) + import importlib + + pipeline = importlib.import_module("pipeline") + loaded = Path(str(pipeline.__file__)).resolve() + if loaded != pipeline_path.resolve(): + raise RuntimeError(f"Imported the wrong ESI-Bench pipeline: {loaded}") + return pipeline + + +def _defer_simulator_shutdown(pipeline: Any) -> Callable[[], None]: + """Keep ``run_one`` alive until the worker response is durable. + + The pinned official runner invokes ``og.shutdown()`` in ``run_one``'s + ``finally`` block. OmniGibson can terminate the interpreter from that call, + before this worker has converted the official result into SHAPER records. + Each worker handles exactly one episode, so deferring shutdown until after + the response file is atomically persisted preserves the official episode + lifecycle without reusing simulator state. + """ + + global _deferred_simulator_shutdown + + omnigibson = getattr(pipeline, "og", None) + shutdown = getattr(omnigibson, "shutdown", None) + if not callable(shutdown): + return lambda: None + + def restore() -> None: + setattr(omnigibson, "shutdown", shutdown) + + def deferred_shutdown() -> Any: + restore() + if getattr(omnigibson, "app", None) is not None: + return shutdown() + return None + + setattr(omnigibson, "shutdown", lambda: None) + _deferred_simulator_shutdown = deferred_shutdown + return restore + + +def _write_response(path: Path, response: Mapping[str, Any]) -> None: + """Atomically persist a complete worker response before simulator exit.""" + + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + payload = json.dumps(response, ensure_ascii=False, indent=2) + "\n" + try: + with temporary.open("w", encoding="utf-8") as stream: + stream.write(payload) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def _finish_deferred_simulator_shutdown() -> None: + """Shut down OmniGibson only after the response is safe to consume.""" + + global _deferred_simulator_shutdown + + shutdown = _deferred_simulator_shutdown + _deferred_simulator_shutdown = None + if shutdown is None: + return + try: + shutdown() + except BaseException: + # Some Isaac/Kit releases terminate the process from shutdown. If they + # instead raise, the already-durable response remains authoritative. + logger.exception("Deferred OmniGibson shutdown raised after response persistence") + + +def _is_environment_failure(exc: BaseException) -> bool: + if isinstance(exc, (PlannerRequestError, HarnessBridgeError)): + return False + messages: list[str] = [] + current: BaseException | None = exc + while current is not None: + messages.append(f"{type(current).__name__}: {current}") + current = current.__cause__ + text = " | ".join(messages).lower() + markers = ( + "omnigibson", + "isaac", + "physx", + "carb", + "cuda", + "vulkan", + "egl", + "renderer", + "render product", + "failed to load usd", + "failed to load scene", + "failed to initialize simulation", + "failed to create render product", + "segmentation fault", + ) + return any(marker in text for marker in markers) + + +def _failure_payload(exc: BaseException) -> tuple[str, bool, str]: + """Classify benchmark failure without turning adapter bugs into zero reward.""" + + if isinstance(exc, PlannerRequestError): + return "planner_failure", False, "planner" + if isinstance(exc, HarnessBridgeError): + return "harness_failure", False, "artifact" + if isinstance(exc, OfficialRunnerError) and _is_environment_failure(exc): + return "environment_failure", True, "environment" + return "adapter_or_upstream_failure", False, "infrastructure" + + +def _history_rounds( + result: Mapping[str, Any], + collector: _HarnessCollector, + task_id: str, +) -> list[dict[str, Any]]: + raw_history = result.get("history") + history: list[dict[str, Any]] = [] + if isinstance(raw_history, list): + for raw_item in cast(list[Any], raw_history): + if not isinstance(raw_item, dict): + raise TypeError("Official ESI-Bench history entries must be dictionaries.") + history.append(cast(dict[str, Any], raw_item)) + image_dir = Path(str(result.get("step_image_dir", ""))) + rounds: list[dict[str, Any]] = [] + for index, item in enumerate(history): + snapshot = collector.snapshots.get(index, {}) + before = snapshot.get("observation_before", []) + after: list[dict[str, Any]] = [] + if index + 1 < len(history): + next_image = image_dir / str(history[index + 1].get("image", "")) + if next_image.is_file(): + after = [collector.images.observation(next_image)["full_frame"]] + for extra in _extra_paths(item): + if extra.is_file(): + after.append(collector.images.observation(extra)["full_frame"]) + if not after and isinstance(before, list): + after = cast(list[dict[str, Any]], list(cast(list[Any], before))) + task_instruction = str(snapshot.get("task_instruction", "")) + if not task_instruction: + task_instruction = "Official ESI-Bench task prompt unavailable for this auxiliary call." + record: dict[str, Any] = { + "record_type": "shaper_round", + "round_index": index, + "task_instruction": task_instruction, + "planner_response": "\n".join( + part + for part in ( + str(item.get("raw_output") or item.get("reasoning", "")).strip(), + str(item.get("raw_output_post_action") or "").strip(), + ) + if part + ), + "command": str(item.get("action", "")), + "observation_before": cast(list[dict[str, Any]], before) if isinstance(before, list) else [], + "observation_after": after, + "context_payload": snapshot.get("context_payload"), + "harness_input": snapshot.get("harness_input"), + "execution_steps": 1, + "action_result": { + **sanitized_action_result(item.get("action_result")), + "action_valid": not bool(item.get("paper_invalid", False)), + "reprompted": bool(item.get("paper_reprompted", False)), + "finish_reason": str(item.get("finish_reason", "")), + }, + "runtime_errors": [], + } + rounds.append(record) + logger.info("Built %d observable rounds for %s", len(rounds), task_id) + return rounds + + +def run_request(request: Mapping[str, Any]) -> dict[str, Any]: + """Run one official question and return only non-privileged SHAPER records.""" + + root = Path(str(request["esi_bench_root"])).expanduser().resolve() + behavior_root = Path(str(request["behavior_root"])).expanduser().resolve() + run_dir = Path(str(request["run_dir"])).expanduser().resolve() + task = cast(dict[str, Any], request["task"]) + planner = cast(dict[str, Any], request["planner"]) + runtime = cast(dict[str, Any], request["runtime"]) + skill = str(request["skill"]) + bridge = cast(dict[str, Any], request["harness_bridge"]) + task_id = str(task.get("task_id", "esi/unknown")) + question_row = load_question_row( + Path(str(request["questions_jsonl"])).expanduser().resolve(), + task.get("question_id"), + ) + expected_runner_task = str(task.get("runner_task", "")).strip() + if str(question_row.get("runner_task", "")).strip() != expected_runner_task: + raise ValueError(f"ESI-Bench runner_task mismatch for {task_id}.") + question_path = resolve_canonical_question( + root / "dataset" / "json_clean", + task.get("question_relpath"), + task.get("question_id"), + ) + canonical_value: object = json.loads(question_path.read_text(encoding="utf-8")) + if not isinstance(canonical_value, dict): + raise ValueError(f"Canonical ESI-Bench question must be an object: {question_path}") + canonical_row = cast(dict[str, Any], canonical_value) + if str(canonical_row.get("runner_task", "")).strip() != expected_runner_task: + raise ValueError(f"Canonical ESI-Bench runner_task mismatch for {task_id}.") + + skill_errors = validate_skill(skill) + if skill_errors: + raise ValueError("Invalid ESI-Bench skill: " + "; ".join(skill_errors)) + build_context = HarnessBridgeClient( + Path(str(bridge["socket_path"])), + str(bridge["token"]), + timeout_seconds=float(bridge["timeout_seconds"]), + max_response_bytes=int(bridge["max_response_bytes"]), + ) + collector = _HarnessCollector( + skill, + build_context, + max_steps=int(runtime["max_steps"]), + ) + model = _OpenAIPlanner( + planner, + skill, + auxiliary_context_builder=collector.route_auxiliary, + ) + pipeline = _load_official_pipeline(root, behavior_root) + restore_shutdown = _defer_simulator_shutdown(pipeline) + original_build_model = pipeline.build_model_client + original_collect = pipeline.collect_contents + + def use_frozen_planner(provider: object, api_key: object, model_name: object) -> _OpenAIPlanner: + del provider, api_key, model_name + return model + + pipeline.build_model_client = use_frozen_planner + pipeline.collect_contents = collector + try: + config = pipeline.ActiveExploreConfig( + task=str(task["runner_task"]), + metadata=question_path, + question_index=0, + json_root=None, + results_root=run_dir / "official_results", + step_image_root=run_dir / "official_steps", + provider="gpt", + model=str(planner["model"]), + api_key=None, + max_steps=int(runtime["max_steps"]), + min_steps=int(runtime["min_steps"]), + threshold=float(runtime["confidence_threshold"]), + max_new_tokens=int(runtime["max_new_tokens"]), + temperature=float(runtime["temperature"]), + top_p=float(runtime["top_p"]), + robot=str(runtime["robot"]), + overwrite=True, + ) + try: + result = cast(dict[str, Any], pipeline.run_one(config)) + except (PlannerRequestError, HarnessBridgeError): + raise + except BaseException as exc: + raise OfficialRunnerError(f"Official ESI-Bench run_one failed: {exc}") from exc + finally: + pipeline.build_model_client = original_build_model + pipeline.collect_contents = original_collect + restore_shutdown() + + skipped = bool(result.get("skipped")) + if skipped: + skip_reason = str(result.get("skip_reason") or "official runner skipped the question") + return { + "ok": False, + "environment_invalid": True, + "termination_reason": "official_skip", + "error": skip_reason, + } + correct = result.get("correct") is True + raw_result_history: object = result.get("history") + final_answer_value: object = result.get("final_answer") + final_answer = cast(dict[str, Any], final_answer_value) if isinstance(final_answer_value, dict) else {} + task_contract = "" + if collector.snapshots: + final_snapshot = collector.snapshots[max(collector.snapshots)] + task_contract = str(final_snapshot.get("task_instruction", "")) + metadata: dict[str, Any] = { + "record_type": "shaper_episode", + "environment_invalid": False, + "termination_reason": str(final_answer.get("stopped_by", "completed")), + "runtime_errors": [], + "extra": { + "task_id": task_id, + "question_id": str(result.get("question_id", task.get("question_id", ""))), + "official_steps": len(cast(list[Any], raw_result_history)) if isinstance(raw_result_history, list) else 0, + "task_family": str(question_row.get("big_task", result.get("task_type", ""))), + "task_subfamily": str(question_row.get("small_task", "")), + "runner_task": expected_runner_task, + "scene": str(result.get("scene", question_row.get("scene", ""))), + "room": str(result.get("room", question_row.get("room", ""))), + "question": str(result.get("question", question_row.get("question", ""))), + "options": _decoded_json_field(question_row, "options_json"), + "task_contract": task_contract, + "final_answer": final_answer, + "ground_truth": result.get("ground_truth", _decoded_json_field(question_row, "answer")), + "official_correct": correct, + "planner_skill": skill, + "reference_images": collector.reference_images, + }, + } + return { + "ok": True, + "reward": 1.0 if correct else 0.0, + "rounds": _history_rounds(result, collector, task_id), + "metadata": metadata, + } + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--response-path", type=Path, required=True) + args = parser.parse_args(argv) + response_path = args.response_path.expanduser().resolve() + response_path.parent.mkdir(parents=True, exist_ok=True) + try: + raw = sys.stdin.read() + request: object = json.loads(raw) + if not isinstance(request, dict): + raise TypeError("Worker request must be a JSON object.") + response = run_request(cast(dict[str, Any], request)) + exit_code = 0 + except BaseException as exc: + termination_reason, environment_invalid, failure_kind = _failure_payload(exc) + response = { + "ok": False, + "environment_invalid": environment_invalid, + "failure_kind": failure_kind, + "termination_reason": termination_reason, + "error": f"{type(exc).__name__}: {exc}", + } + logger.exception("ESI-Bench worker failed") + exit_code = 2 + _write_response(response_path, response) + _finish_deferred_simulator_shutdown() + return exit_code + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/contrib/recipes/shaper/evaluate.py b/contrib/recipes/shaper/evaluate.py new file mode 100644 index 000000000..b3d990331 --- /dev/null +++ b/contrib/recipes/shaper/evaluate.py @@ -0,0 +1,216 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Evaluate one SHAPER artifact pair on a benchmark split.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence, cast + +import agentlightning as agl +from agentlightning.algorithm import Baseline +from agentlightning.types import Dataset, PromptTemplate +from contrib.agentlightning.contrib.shaper import EpisodeTrace, SHAPERTraceAdapter + +from .reproduce import ReproductionBundle, execution_strategy, load_bundle_factory + + +@dataclass(frozen=True) +class EvaluationConfig: + """Configuration for an official simulator evaluation.""" + + factory: str + split: str + start_index: int + limit: int | None + n_runners: int + output: Path + skill_path: Path | None + harness_path: Path | None + + +def parse_args(argv: Sequence[str] | None = None) -> EvaluationConfig: + parser = argparse.ArgumentParser(description="Evaluate SHAPER artifacts on a benchmark split.") + parser.add_argument("--factory", required=True, help="Benchmark factory as module:function.") + parser.add_argument("--split", choices=("train", "validation"), default="validation") + parser.add_argument("--start-index", type=int, default=0) + parser.add_argument("--limit", type=int, help="Number of episodes; omit to evaluate the rest of the split.") + parser.add_argument("--n-runners", type=int, default=1) + parser.add_argument("--output", type=Path, default=Path("outputs/shaper/evaluation.json")) + parser.add_argument("--skill-path", type=Path, help="Skill artifact; defaults to the benchmark seed.") + parser.add_argument("--harness-path", type=Path, help="Harness artifact; defaults to the benchmark seed.") + args = parser.parse_args(argv) + return EvaluationConfig( + factory=str(args.factory), + split=str(args.split), + start_index=int(args.start_index), + limit=cast(int | None, args.limit), + n_runners=int(args.n_runners), + output=cast(Path, args.output), + skill_path=cast(Path | None, args.skill_path), + harness_path=cast(Path | None, args.harness_path), + ) + + +def select_tasks(dataset: Dataset[Any], *, start_index: int, limit: int | None) -> list[Any]: + """Select a contiguous, non-empty evaluation range.""" + + if start_index < 0: + raise ValueError("start-index must not be negative.") + if limit is not None and limit < 1: + raise ValueError("limit must be positive when provided.") + if start_index >= len(dataset): + raise IndexError(f"start-index {start_index} is outside a dataset of size {len(dataset)}.") + stop = len(dataset) if limit is None else min(len(dataset), start_index + limit) + return [dataset[index] for index in range(start_index, stop)] + + +def _resource_override(path: Path | None, fallback: PromptTemplate) -> PromptTemplate: + if path is None: + return fallback + resolved = path.expanduser().resolve() + return PromptTemplate(template=resolved.read_text(encoding="utf-8"), engine="f-string") + + +def prepare_resources( + bundle: ReproductionBundle[Any], + *, + skill_path: Path | None, + harness_path: Path | None, +) -> dict[str, Any]: + """Load and validate the artifact pair used for evaluation.""" + + resources = dict(bundle.initial_resources) + seed_skill = resources.get(bundle.skill_resource_name) + seed_harness = resources.get(bundle.harness_resource_name) + if not isinstance(seed_skill, PromptTemplate) or not isinstance(seed_harness, PromptTemplate): + raise TypeError("Benchmark factories must provide PromptTemplate skill and harness resources.") + skill = _resource_override(skill_path, seed_skill) + harness = _resource_override(harness_path, seed_harness) + skill_errors = list(bundle.skill_validator(skill.template)) + if skill_errors: + raise ValueError("Skill failed benchmark validation: " + "; ".join(skill_errors)) + harness_result = bundle.harness_validator.validate(harness.template) + if not harness_result.valid: + raise ValueError("Harness failed benchmark validation: " + "; ".join(harness_result.errors)) + resources[bundle.skill_resource_name] = skill + resources[bundle.harness_resource_name] = harness + return resources + + +def _safe_task_identity(task: object) -> dict[str, Any]: + """Return identifiers only, never simulator configuration or scorer labels.""" + + if not isinstance(task, Mapping): + return {"type": type(task).__name__} + allowed = ("task_id", "task_name", "question_id", "runner_task", "episode_index", "max_steps") + return {key: task[key] for key in allowed if key in task} + + +async def _collect_episodes(trainer: agl.Trainer, expected: int) -> list[dict[str, Any]]: + rollouts = list(await trainer.store.query_rollouts(sort_by="start_time", sort_order="asc")) + if len(rollouts) != expected: + raise RuntimeError(f"Evaluation expected {expected} rollouts, found {len(rollouts)}.") + + episodes: list[dict[str, Any]] = [] + for rollout in rollouts: + spans = await trainer.store.query_spans(rollout_id=rollout.rollout_id, attempt_id="latest") + trace: EpisodeTrace = ( + SHAPERTraceAdapter() + .adapt(spans) + .model_copy(update={"rollout_id": rollout.rollout_id, "task": None, "status": rollout.status}) + ) + episodes.append( + { + "rollout_id": rollout.rollout_id, + "task": _safe_task_identity(rollout.input), + "status": rollout.status, + "reward": trace.final_reward, + "environment_invalid": trace.metadata.environment_invalid, + "termination_reason": trace.metadata.termination_reason, + "round_count": len(trace.rounds), + "runtime_errors": [ + *trace.metadata.runtime_errors, + *(error for record in trace.rounds for error in record.runtime_errors), + ], + "adapter_errors": trace.adapter_errors, + } + ) + return episodes + + +def _summarize(episodes: Sequence[Mapping[str, Any]]) -> dict[str, Any]: + rewards = [float(item["reward"]) for item in episodes if isinstance(item.get("reward"), (int, float))] + invalid_count = sum(bool(item.get("environment_invalid")) for item in episodes) + failed_count = sum(item.get("status") != "succeeded" for item in episodes) + adapter_error_count = sum(bool(item.get("adapter_errors")) for item in episodes) + evaluation_valid = ( + len(rewards) == len(episodes) and invalid_count == 0 and failed_count == 0 and adapter_error_count == 0 + ) + reward_sum = sum(rewards) + return { + "episode_count": len(episodes), + "scored_episodes": len(rewards), + "reward_sum": reward_sum, + "mean_reward": reward_sum / len(rewards) if rewards else None, + "environment_invalid_episodes": invalid_count, + "failed_rollouts": failed_count, + "adapter_error_episodes": adapter_error_count, + "evaluation_valid": evaluation_valid, + } + + +def run_evaluation(config: EvaluationConfig) -> dict[str, Any]: + """Run official benchmark rollouts and persist aggregate and per-episode results.""" + + if config.n_runners < 1: + raise ValueError("n-runners must be at least 1.") + bundle = load_bundle_factory(config.factory)() + dataset = bundle.train_dataset if config.split == "train" else bundle.val_dataset + tasks = select_tasks(dataset, start_index=config.start_index, limit=config.limit) + resources = prepare_resources( + bundle, + skill_path=config.skill_path, + harness_path=config.harness_path, + ) + trainer = agl.Trainer( + algorithm=Baseline( + polling_interval=0.05, + max_queue_length=max(1, config.n_runners * 2), + span_verbosity="none", + ), + adapter=SHAPERTraceAdapter(), + initial_resources=resources, + n_runners=config.n_runners, + strategy=execution_strategy(config.n_runners), + tracer=agl.OtelTracer(), + ) + trainer.dev(agent=bundle.agent, train_dataset=tasks) + episodes = asyncio.run(_collect_episodes(trainer, len(tasks))) + report = { + "factory": config.factory, + "split": config.split, + "start_index": config.start_index, + "skill_path": str(config.skill_path.expanduser().resolve()) if config.skill_path else None, + "harness_path": str(config.harness_path.expanduser().resolve()) if config.harness_path else None, + "summary": _summarize(episodes), + "episodes": episodes, + } + output = config.output.expanduser().resolve() + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + return report + + +def main(argv: Sequence[str] | None = None) -> int: + report = run_evaluation(parse_args(argv)) + print(json.dumps(report["summary"], ensure_ascii=False, indent=2)) + return 0 if report["summary"]["evaluation_valid"] else 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/contrib/recipes/shaper/harness_bridge.py b/contrib/recipes/shaper/harness_bridge.py new file mode 100644 index 000000000..2999ed713 --- /dev/null +++ b/contrib/recipes/shaper/harness_bridge.py @@ -0,0 +1,183 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Private JSON bridge between simulator workers and the harness sandbox.""" + +from __future__ import annotations + +import json +import os +import secrets +import socket +import struct +import tempfile +import threading +from pathlib import Path +from typing import Any, Callable, Mapping, cast + +_HEADER = struct.Struct("!Q") + + +class HarnessBridgeError(RuntimeError): + """Raised when a harness bridge request cannot be completed.""" + + +def _receive_exact(connection: socket.socket, size: int) -> bytes: + chunks: list[bytes] = [] + remaining = size + while remaining: + chunk = connection.recv(remaining) + if not chunk: + raise HarnessBridgeError("Harness bridge connection closed before the message completed.") + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) + + +def _receive_json(connection: socket.socket, max_bytes: int) -> dict[str, Any]: + raw_size = _receive_exact(connection, _HEADER.size) + size = _HEADER.unpack(raw_size)[0] + if size > max_bytes: + raise HarnessBridgeError(f"Harness bridge message is {size} bytes; limit is {max_bytes} bytes.") + raw = _receive_exact(connection, size) + try: + value: object = json.loads(raw.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise HarnessBridgeError(f"Harness bridge received invalid JSON: {exc}") from exc + if not isinstance(value, dict): + raise HarnessBridgeError("Harness bridge message must be a JSON object.") + return cast(dict[str, Any], value) + + +def _send_json(connection: socket.socket, value: Mapping[str, Any], max_bytes: int) -> None: + raw = json.dumps(dict(value), ensure_ascii=True, separators=(",", ":")).encode("utf-8") + if len(raw) > max_bytes: + raise HarnessBridgeError(f"Harness bridge response is {len(raw)} bytes; limit is {max_bytes} bytes.") + connection.sendall(_HEADER.pack(len(raw)) + raw) + + +class HarnessBridgeClient: + """Call a controller-owned harness runtime from an isolated worker.""" + + def __init__( + self, + socket_path: Path, + token: str, + *, + timeout_seconds: float = 10.0, + max_request_bytes: int = 256_000_000, + max_response_bytes: int = 32_000_000, + ) -> None: + self.socket_path = socket_path + self.token = token + self.timeout_seconds = timeout_seconds + self.max_request_bytes = max_request_bytes + self.max_response_bytes = max_response_bytes + + def __call__(self, records: list[dict[str, Any]]) -> Any: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.settimeout(self.timeout_seconds) + try: + connection.connect(str(self.socket_path)) + _send_json( + connection, + {"token": self.token, "records": records}, + self.max_request_bytes, + ) + response = _receive_json(connection, self.max_response_bytes) + except (OSError, TimeoutError) as exc: + raise HarnessBridgeError(f"Harness bridge request failed: {exc}") from exc + if response.get("ok") is not True: + raise HarnessBridgeError(str(response.get("error", "Harness bridge rejected the request."))) + return response.get("output") + + +class HarnessBridgeServer: + """Serve one restricted harness runtime on a private Unix socket.""" + + def __init__( + self, + handler: Callable[[list[dict[str, Any]]], Any], + *, + max_request_bytes: int = 256_000_000, + max_response_bytes: int = 32_000_000, + ) -> None: + self.handler = handler + self.max_request_bytes = max_request_bytes + self.max_response_bytes = max_response_bytes + self.token = secrets.token_hex(32) + self.socket_path: Path | None = None + self._directory: tempfile.TemporaryDirectory[str] | None = None + self._listener: socket.socket | None = None + self._thread: threading.Thread | None = None + self._stop = threading.Event() + self._ready = threading.Event() + self._startup_error: BaseException | None = None + + def __enter__(self) -> HarnessBridgeServer: + self._directory = tempfile.TemporaryDirectory(prefix="shaper-harness-") + self.socket_path = Path(self._directory.name) / "bridge.sock" + self._thread = threading.Thread(target=self._serve, name="shaper-harness-bridge", daemon=True) + self._thread.start() + if not self._ready.wait(timeout=5.0): + self.close() + raise HarnessBridgeError("Harness bridge did not become ready within five seconds.") + if self._startup_error is not None: + error = self._startup_error + self.close() + raise HarnessBridgeError(f"Harness bridge failed to start: {error}") from error + return self + + def _serve(self) -> None: + assert self.socket_path is not None + try: + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + self._listener = listener + listener.bind(str(self.socket_path)) + os.chmod(self.socket_path, 0o600) + listener.listen(1) + listener.settimeout(0.2) + except BaseException as exc: + self._startup_error = exc + self._ready.set() + return + self._ready.set() + while not self._stop.is_set(): + try: + connection, _ = listener.accept() + except socket.timeout: + continue + except OSError: + break + with connection: + connection.settimeout(15.0) + try: + request = _receive_json(connection, self.max_request_bytes) + if not secrets.compare_digest(str(request.get("token", "")), self.token): + raise HarnessBridgeError("Harness bridge authentication failed.") + raw_records = request.get("records") + if not isinstance(raw_records, list) or not all(isinstance(item, dict) for item in raw_records): + raise HarnessBridgeError("Harness bridge records must be a list of JSON objects.") + records = cast(list[dict[str, Any]], raw_records) + response: dict[str, Any] = {"ok": True, "output": self.handler(records)} + except BaseException as exc: + response = {"ok": False, "error": f"{type(exc).__name__}: {exc}"} + try: + _send_json(connection, response, self.max_response_bytes) + except (HarnessBridgeError, OSError): + pass + + def close(self) -> None: + self._stop.set() + if self._listener is not None: + self._listener.close() + if self._thread is not None: + self._thread.join(timeout=2.0) + if self._directory is not None: + self._directory.cleanup() + self._listener = None + self._thread = None + self._directory = None + + def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None: + del exc_type, exc_value, traceback + self.close() diff --git a/contrib/recipes/shaper/integration.py b/contrib/recipes/shaper/integration.py new file mode 100644 index 000000000..c0af3b347 --- /dev/null +++ b/contrib/recipes/shaper/integration.py @@ -0,0 +1,39 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Small helpers for connecting an embodied agent to the SHAPER trace contract.""" + +from __future__ import annotations + +from typing import Any, Callable, Dict + +from agentlightning.types import NamedResources, PromptTemplate +from contrib.agentlightning.contrib.shaper import PythonHarnessValidator + + +def get_artifact_text( + resources: NamedResources, + *, + skill_resource_name: str = "skill", + harness_resource_name: str = "harness", +) -> tuple[str, str]: + """Extract skill text and harness source from an AGL resource bundle.""" + + skill = resources.get(skill_resource_name) + harness = resources.get(harness_resource_name) + if not isinstance(skill, PromptTemplate): + raise TypeError(f"{skill_resource_name!r} must be a PromptTemplate resource.") + if not isinstance(harness, PromptTemplate): + raise TypeError(f"{harness_resource_name!r} must be a PromptTemplate resource.") + return skill.template, harness.template + + +def load_context_builder( + source: str, + *, + validator: PythonHarnessValidator | None = None, +) -> Callable[[list[Dict[str, Any]]], Any]: + """Validate a context builder and return its restricted-process callable.""" + + effective_validator = validator or PythonHarnessValidator() + runtime = effective_validator.runtime(source) + return runtime diff --git a/contrib/recipes/shaper/pyrightconfig.json b/contrib/recipes/shaper/pyrightconfig.json new file mode 100644 index 000000000..2ac3d884e --- /dev/null +++ b/contrib/recipes/shaper/pyrightconfig.json @@ -0,0 +1,23 @@ +{ + "include": [ + "../../agentlightning/contrib/shaper", + "integration.py", + "reproduce.py", + "evaluate.py", + "common.py", + "vlabench", + "esi_bench" + ], + "exclude": [ + "**/__pycache__", + "vlabench/prompts/*.py", + "esi_bench/prompts/*.py" + ], + "extraPaths": ["../../.."], + "venvPath": "../../..", + "venv": ".venv", + "pythonVersion": "3.12", + "typeCheckingMode": "strict", + "reportMissingTypeStubs": "none", + "reportMissingModuleSource": "none" +} diff --git a/contrib/recipes/shaper/reproduce.py b/contrib/recipes/shaper/reproduce.py new file mode 100644 index 000000000..25d987862 --- /dev/null +++ b/contrib/recipes/shaper/reproduce.py @@ -0,0 +1,350 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Run SHAPER against an Agent Lightning benchmark bundle. + +Usage from the repository root: + + OPENAI_API_KEY=... python -m contrib.recipes.shaper.reproduce \ + --factory contrib.recipes.shaper.vlabench.factory:build_bundle \ + --model qwen3.6-27b \ + --output-dir outputs/shaper + +Complete factories are included for VLABench and ESI-Bench; external adapters +may provide the same :class:`ReproductionBundle` contract. This module never +stores API keys in its output. +""" + +from __future__ import annotations + +import argparse +import importlib +import json +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Generic, Mapping, Sequence, TypeVar, cast + +from openai import AsyncOpenAI + +import agentlightning as agl +from agentlightning.litagent import LitAgent +from agentlightning.types import LLM, Dataset, NamedResources, ProxyLLM +from contrib.agentlightning.contrib.shaper import ( + DEFAULT_HARNESS_CONTRACT, + SHAPER, + PythonHarnessValidator, + SHAPERRoleProtocol, + SHAPERTraceAdapter, + SkillValidator, + validate_nonempty_skill, +) + +T_task = TypeVar("T_task") + + +@dataclass(frozen=True) +class ReproductionBundle(Generic[T_task]): + """Benchmark objects required by the generic reproduction runner. + + ``planner_resource_name`` identifies the actual LLM resource consumed by + the embodied planner. The optimizer reuses that resource's model and + endpoint; factory-declared strings are not trusted as evidence of model + identity. + """ + + agent: LitAgent[T_task] + train_dataset: Dataset[T_task] + val_dataset: Dataset[T_task] + initial_resources: NamedResources + planner_resource_name: str = "planner_llm" + skill_resource_name: str = "skill" + harness_resource_name: str = "harness" + harness_contract: str = DEFAULT_HARNESS_CONTRACT + skill_validator: SkillValidator = validate_nonempty_skill + harness_validator: PythonHarnessValidator = field(default_factory=PythonHarnessValidator) + judger_prompt: str | None = None + summarizer_prompt: str | None = None + skill_optimizer_prompt: str | None = None + harness_optimizer_prompt: str | None = None + role_protocol: SHAPERRoleProtocol | None = None + provenance: Mapping[str, Any] = field(default_factory=lambda: cast(dict[str, Any], {})) + + +@dataclass(frozen=True) +class ReproductionConfig: + """Typed command-line configuration with the SHAPER recipe defaults.""" + + factory: str + model: str | None + output_dir: Path + base_url: str | None + api_key_env: str + n_runners: int + validation_size: int | None + gradient_batch_size: int + beam_width: int + branch_factor: int + skill_rounds: int + harness_rounds: int + role_max_completion_tokens: int | None + optimizer_temperature: float + rollout_batch_timeout: float + artifact_repair_attempts: int + random_seed: int + + +BundleFactory = Callable[[], ReproductionBundle[Any]] + +_SENSITIVE_PROVENANCE_KEYS = frozenset( + { + "api_key", + "authorization", + "credential", + "credentials", + "password", + "secret", + "token", + } +) + + +def validate_provenance(value: Mapping[str, Any]) -> dict[str, Any]: + """Return a JSON-safe provenance object and reject likely credentials. + + Reproduction metadata is intentionally public-facing. Factories must record + immutable source and protocol identifiers, never endpoint credentials. + """ + + def reject_sensitive_keys(item: object, path: str) -> None: + if isinstance(item, Mapping): + mapping = cast(Mapping[object, object], item) + for raw_key, child in mapping.items(): + if not isinstance(raw_key, str): + raise TypeError(f"Provenance key at {path} must be a string.") + normalized = raw_key.lower().replace("-", "_") + if normalized in _SENSITIVE_PROVENANCE_KEYS or normalized.endswith( + ("_api_key", "_authorization", "_credential", "_credentials", "_password", "_secret") + ): + raise ValueError(f"Sensitive provenance key is not allowed: {path}.{raw_key}") + reject_sensitive_keys(child, f"{path}.{raw_key}") + elif isinstance(item, (list, tuple)): + sequence = cast(Sequence[object], item) + for index, child in enumerate(sequence): + reject_sensitive_keys(child, f"{path}[{index}]") + + copied = dict(value) + reject_sensitive_keys(copied, "provenance") + try: + encoded = json.dumps(copied, ensure_ascii=False, allow_nan=False) + except (TypeError, ValueError) as exc: + raise TypeError(f"Reproduction provenance must be finite JSON data: {exc}") from exc + decoded: object = json.loads(encoded) + if not isinstance(decoded, dict): + raise TypeError("Reproduction provenance must be a JSON object.") + return cast(dict[str, Any], decoded) + + +def load_bundle_factory(spec: str) -> BundleFactory: + """Load a ``module:function`` bundle factory without importing benchmark code here.""" + + module_name, separator, attribute_name = spec.partition(":") + if not separator or not module_name or not attribute_name: + raise ValueError("Factory must use the form 'module:function'.") + module = importlib.import_module(module_name) + value: object = getattr(module, attribute_name) + if not callable(value): + raise TypeError(f"Factory {spec!r} is not callable.") + return cast(BundleFactory, value) + + +def execution_strategy(n_runners: int) -> str | dict[str, object]: + """Run every simulator worker in its own process.""" + + if n_runners < 1: + raise ValueError("n_runners must be at least 1.") + return {"type": "cs", "main_process": "algorithm"} + + +def parse_args(argv: Sequence[str] | None = None) -> ReproductionConfig: + """Parse CLI flags into a type-checked immutable configuration.""" + + parser = argparse.ArgumentParser(description="Run two-stage SHAPER artifact evolution.") + parser.add_argument("--factory", required=True, help="Benchmark factory as module:function.") + parser.add_argument( + "--model", + help="Optional assertion for the model declared by the planner LLM resource.", + ) + parser.add_argument("--output-dir", type=Path, default=Path("outputs/shaper")) + parser.add_argument("--base-url", help="Optional assertion for the planner endpoint declared by the factory.") + parser.add_argument("--api-key-env", default="OPENAI_API_KEY") + parser.add_argument("--n-runners", type=int, default=1) + parser.add_argument("--validation-size", type=int) + parser.add_argument("--gradient-batch-size", type=int, default=4) + parser.add_argument("--beam-width", type=int, default=3) + parser.add_argument("--branch-factor", type=int, default=2) + parser.add_argument("--skill-rounds", type=int, default=2) + parser.add_argument("--harness-rounds", type=int, default=2) + parser.add_argument( + "--role-max-completion-tokens", + type=int, + help="Override the optimizer-role output limit declared by the planner resource.", + ) + parser.add_argument("--optimizer-temperature", type=float, default=1.0) + parser.add_argument( + "--rollout-batch-timeout", + type=float, + default=3600.0, + help="Wall-clock allowance per concurrent rollout wave (default: 3600s).", + ) + parser.add_argument("--artifact-repair-attempts", type=int, default=1) + parser.add_argument("--random-seed", type=int, default=0) + args = parser.parse_args(argv) + return ReproductionConfig( + factory=str(args.factory), + model=cast(str | None, args.model), + output_dir=cast(Path, args.output_dir), + base_url=cast(str | None, args.base_url), + api_key_env=str(args.api_key_env), + n_runners=int(args.n_runners), + validation_size=cast(int | None, args.validation_size), + gradient_batch_size=int(args.gradient_batch_size), + beam_width=int(args.beam_width), + branch_factor=int(args.branch_factor), + skill_rounds=int(args.skill_rounds), + harness_rounds=int(args.harness_rounds), + role_max_completion_tokens=cast(int | None, args.role_max_completion_tokens), + optimizer_temperature=float(args.optimizer_temperature), + rollout_batch_timeout=float(args.rollout_batch_timeout), + artifact_repair_attempts=int(args.artifact_repair_attempts), + random_seed=int(args.random_seed), + ) + + +def run_reproduction(config: ReproductionConfig) -> SHAPER[Any]: + """Run one benchmark bundle and persist JSON-serializable optimization artifacts.""" + + if config.n_runners < 1: + raise ValueError("n_runners must be at least 1.") + bundle = load_bundle_factory(config.factory)() + provenance = validate_provenance(bundle.provenance) + planner = bundle.initial_resources.get(bundle.planner_resource_name) + if not isinstance(planner, LLM): + raise TypeError(f"Resource {bundle.planner_resource_name!r} must be an LLM used by the embodied planner.") + if config.model is not None and planner.model != config.model: + raise ValueError( + "SHAPER reproduction requires one model identity for planner and optimizer: " + f"planner resource uses {planner.model!r}, CLI asserted {config.model!r}." + ) + planner_endpoint = planner.get_base_url(None, None) if isinstance(planner, ProxyLLM) else planner.get_base_url() + if config.base_url is not None and config.base_url.rstrip("/") != planner_endpoint.rstrip("/"): + raise ValueError( + "The optimizer endpoint cannot differ from the embodied planner endpoint: " + f"planner uses {planner_endpoint!r}, CLI asserted {config.base_url!r}." + ) + # Local OpenAI-compatible servers commonly do not authenticate. The SDK + # still requires a non-empty string, while remote providers will reject the + # placeholder clearly if the user forgot their real credential. + api_key = os.environ.get(config.api_key_env) or planner.api_key or "not-required" + + sampling: dict[str, Any] = planner.sampling_parameters + client = AsyncOpenAI( + api_key=api_key, + base_url=planner_endpoint, + timeout=float(sampling.get("timeout", 300.0)), + max_retries=int(sampling.get("max_retries", 2)), + ) + role_max_completion_tokens = config.role_max_completion_tokens or int( + sampling.get("optimizer_max_completion_tokens", 65_536) + ) + raw_extra_body: object = sampling.get("extra_body") + if raw_extra_body is not None and not isinstance(raw_extra_body, dict): + raise TypeError("planner sampling_parameters.extra_body must be a dictionary when provided.") + role_extra_body = dict(cast(dict[str, Any], raw_extra_body)) if isinstance(raw_extra_body, dict) else {} + for key in ("top_p", "presence_penalty"): + if key in sampling: + role_extra_body.setdefault(key, sampling[key]) + algorithm = SHAPER[Any]( + client, + model=planner.model, + skill_resource_name=bundle.skill_resource_name, + harness_resource_name=bundle.harness_resource_name, + validation_size=config.validation_size, + gradient_batch_size=config.gradient_batch_size, + beam_width=config.beam_width, + branch_factor=config.branch_factor, + skill_rounds=config.skill_rounds, + harness_rounds=config.harness_rounds, + rollout_batch_timeout=config.rollout_batch_timeout, + optimizer_temperature=config.optimizer_temperature, + role_max_completion_tokens=role_max_completion_tokens, + role_extra_body=role_extra_body or None, + artifact_repair_attempts=config.artifact_repair_attempts, + random_seed=config.random_seed, + harness_contract=bundle.harness_contract, + skill_validator=bundle.skill_validator, + harness_validator=bundle.harness_validator, + judger_prompt=bundle.judger_prompt, + summarizer_prompt=bundle.summarizer_prompt, + skill_optimizer_prompt=bundle.skill_optimizer_prompt, + harness_optimizer_prompt=bundle.harness_optimizer_prompt, + role_protocol=bundle.role_protocol, + ) + trainer = agl.Trainer( + algorithm=algorithm, + adapter=SHAPERTraceAdapter(), + initial_resources=bundle.initial_resources, + n_runners=config.n_runners, + strategy=execution_strategy(config.n_runners), + tracer=agl.OtelTracer(), + ) + trainer.fit( + agent=bundle.agent, + train_dataset=bundle.train_dataset, + val_dataset=bundle.val_dataset, + ) + + config.output_dir.mkdir(parents=True, exist_ok=True) + best = algorithm.get_best_candidate() + report = { + "configuration": { + "factory": config.factory, + "model": planner.model, + "planner_resource_name": bundle.planner_resource_name, + "n_runners": config.n_runners, + "planner_endpoint_configured": bool(planner_endpoint), + "validation_size": config.validation_size, + "gradient_batch_size": config.gradient_batch_size, + "beam_width": config.beam_width, + "branch_factor": config.branch_factor, + "skill_rounds": config.skill_rounds, + "harness_rounds": config.harness_rounds, + "role_max_completion_tokens": role_max_completion_tokens, + "optimizer_temperature": config.optimizer_temperature, + "rollout_batch_timeout": config.rollout_batch_timeout, + "artifact_repair_attempts": config.artifact_repair_attempts, + "provider_extra_body_configured": bool(role_extra_body), + "random_seed": config.random_seed, + }, + "provenance": provenance, + "best_candidate": best.model_dump(mode="json"), + "optimization_history": [event.model_dump(mode="json") for event in algorithm.get_optimization_history()], + } + (config.output_dir / "shaper_run.json").write_text( + json.dumps(report, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8", + ) + (config.output_dir / "best_skill.txt").write_text(best.skill.template.rstrip() + "\n", encoding="utf-8") + (config.output_dir / "best_harness.py").write_text(best.harness.template.rstrip() + "\n", encoding="utf-8") + return algorithm + + +def main(argv: Sequence[str] | None = None) -> None: + """CLI entry point.""" + + algorithm = run_reproduction(parse_args(argv)) + best = algorithm.get_best_candidate() + print(f"SHAPER complete: {best.version} validation_score={best.validation_score:.4f}") + + +if __name__ == "__main__": + main() diff --git a/contrib/recipes/shaper/scripts/bootstrap_shaper_environment.sh b/contrib/recipes/shaper/scripts/bootstrap_shaper_environment.sh new file mode 100755 index 000000000..2e3b01440 --- /dev/null +++ b/contrib/recipes/shaper/scripts/bootstrap_shaper_environment.sh @@ -0,0 +1,246 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft. All rights reserved. + +set -euo pipefail + +VLABENCH_COMMIT="cf588fe60c0c7282174fe979f5913170cfe69017" +OPENPI_COMMIT="4483d1da6332da44115fe530e4e6fdd89bd57b13" +ESI_BENCH_COMMIT="3c1756396f32b1a90c1f72356a7fde45f418e179" +BEHAVIOR_COMMIT="67ad490856dd465d4606663106f81673fc8bf4e8" +GOOGLE_GENAI_VERSION="1.75.0" +FASTAPI_VERSION="0.121.2" +STARLETTE_VERSION="0.49.3" + +usage() { + cat <<'EOF' +Usage: + bootstrap_shaper_environment.sh common AGL_ROOT + bootstrap_shaper_environment.sh vlabench-simulator AGL_ROOT VLABENCH_CHECKOUT OPENPI_ROOT [--download-assets] + bootstrap_shaper_environment.sh vlabench-actor AGL_ROOT OPENPI_ROOT CHECKPOINT_DIR [--download-checkpoint] + bootstrap_shaper_environment.sh esi-controller AGL_ROOT + bootstrap_shaper_environment.sh esi-worker AGL_ROOT ESI_ROOT BEHAVIOR_ROOT [--install-behavior] + +Run each mode inside its intended Python environment. The ESI worker mode +requires an active Python 3.11 conda environment plus explicit acceptance: + SHAPER_ACCEPT_NVIDIA_EULA=YES + SHAPER_ACCEPT_BEHAVIOR_DATASET_TOS=YES +and requires OMNIGIBSON_DATA_PATH to point at a data disk. +EOF +} + +die() { + echo "[shaper-bootstrap] $*" >&2 + exit 2 +} + +python_bin() { + if [[ -n "${PYTHON:-}" ]]; then + printf '%s\n' "$PYTHON" + elif command -v python >/dev/null 2>&1; then + command -v python + elif command -v python3 >/dev/null 2>&1; then + command -v python3 + else + die "No Python interpreter found. Set PYTHON to the target environment interpreter." + fi +} + +absolute_dir() { + local path="$1" + [[ -d "$path" ]] || die "Missing directory: $path" + (cd "$path" && pwd) +} + +require_commit() { + local checkout="$1" + local expected="$2" + local label="$3" + local actual + actual="$(git -C "$checkout" rev-parse HEAD 2>/dev/null)" || die "$label is not a Git checkout: $checkout" + [[ "$actual" == "$expected" ]] || die "$label revision $actual does not match pinned $expected" +} + +require_python_version() { + local executable="$1" + local expected="$2" + local actual + actual="$($executable -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')" + [[ "$actual" == "$expected" ]] || die "Expected Python $expected, found $actual at $executable" +} + +install_agl() { + local executable="$1" + local agl_root="$2" + "$executable" -m pip install -e "$agl_root" + # Editable installs resolve the root dependency set and can otherwise + # upgrade these packages again. Pin them last to the versions validated by + # the repository lock; newer FastAPI releases remove APIs used by LiteLLM. + "$executable" -m pip install --upgrade --force-reinstall \ + "fastapi==$FASTAPI_VERSION" "starlette==$STARLETTE_VERSION" + (cd "$agl_root" && "$executable" -c 'from contrib.agentlightning.contrib.shaper import SHAPER; print("SHAPER import:", SHAPER.__name__)') +} + +mode="${1:-}" +if [[ -z "$mode" ]]; then + usage + exit 2 +fi +shift + +case "$mode" in + common) + [[ "$#" -eq 1 ]] || { usage; exit 2; } + agl_root="$(absolute_dir "$1")" + py="$(python_bin)" + install_agl "$py" "$agl_root" + ;; + + vlabench-simulator) + [[ "$#" -ge 3 && "$#" -le 4 ]] || { usage; exit 2; } + agl_root="$(absolute_dir "$1")" + vlabench_checkout="$(absolute_dir "$2")" + openpi_root="$(absolute_dir "$3")" + option="${4:-}" + [[ -z "$option" || "$option" == "--download-assets" ]] || die "Unknown option: $option" + require_commit "$vlabench_checkout" "$VLABENCH_COMMIT" "VLABench" + require_commit "$openpi_root" "$OPENPI_COMMIT" "OpenPI" + py="$(python_bin)" + require_python_version "$py" "3.10" + export VLABENCH_ROOT="$vlabench_checkout/VLABench" + "$py" -m pip install -r "$agl_root/contrib/recipes/shaper/vlabench/requirements-simulator.txt" + "$py" -m pip install --no-deps -e "$vlabench_checkout" + "$py" -m pip install -e "$openpi_root/packages/openpi-client" + install_agl "$py" "$agl_root" + if [[ "$option" == "--download-assets" ]]; then + (cd "$vlabench_checkout" && "$py" scripts/download_assets.py --choice all) + fi + PYTHONPATH="$agl_root${PYTHONPATH:+:$PYTHONPATH}" "$py" - "$VLABENCH_ROOT" <<'PY' +from pathlib import Path +import sys + +from contrib.recipes.shaper.vlabench.check_env import check_vlabench_assets + +errors = check_vlabench_assets(Path(sys.argv[1])) +if errors: + raise SystemExit("\n".join(errors) + "\nRe-run with --download-assets.") +import VLABench # noqa: F401 +import openpi_client # noqa: F401 +print("VLABench simulator imports and assets passed.") +PY + ;; + + vlabench-actor) + [[ "$#" -ge 3 && "$#" -le 4 ]] || { usage; exit 2; } + agl_root="$(absolute_dir "$1")" + openpi_root="$(absolute_dir "$2")" + checkpoint_dir="$3" + option="${4:-}" + [[ -z "$option" || "$option" == "--download-checkpoint" ]] || die "Unknown option: $option" + require_commit "$openpi_root" "$OPENPI_COMMIT" "OpenPI" + command -v uv >/dev/null 2>&1 || die "uv is required for the pinned OpenPI environment." + (cd "$openpi_root" && GIT_LFS_SKIP_SMUDGE=1 uv sync --frozen --no-dev) + actor_python="$openpi_root/.venv/bin/python" + [[ -x "$actor_python" ]] || die "OpenPI uv environment did not create $actor_python" + if [[ "$option" == "--download-checkpoint" ]]; then + PYTHONPATH="$agl_root${PYTHONPATH:+:$PYTHONPATH}" \ + "$actor_python" "$agl_root/contrib/recipes/shaper/scripts/download_shaper_vlabench_actor.py" \ + "$checkpoint_dir" + else + PYTHONPATH="$agl_root${PYTHONPATH:+:$PYTHONPATH}" \ + "$actor_python" "$agl_root/contrib/recipes/shaper/scripts/download_shaper_vlabench_actor.py" \ + "$checkpoint_dir" --verify-only + fi + PYTHONPATH="$agl_root${PYTHONPATH:+:$PYTHONPATH}" "$actor_python" -c \ + 'import openpi; import contrib.recipes.shaper.vlabench.openpi_server; print("OpenPI actor imports passed.")' + ;; + + esi-controller) + [[ "$#" -eq 1 ]] || { usage; exit 2; } + agl_root="$(absolute_dir "$1")" + py="$(python_bin)" + install_agl "$py" "$agl_root" + ;; + + esi-worker) + [[ "$#" -ge 3 && "$#" -le 4 ]] || { usage; exit 2; } + agl_root="$(absolute_dir "$1")" + esi_root="$(absolute_dir "$2")" + behavior_root="$(absolute_dir "$3")" + option="${4:-}" + [[ -z "$option" || "$option" == "--install-behavior" ]] || die "Unknown option: $option" + [[ "$(uname -s)" == "Linux" && "$(uname -m)" == "x86_64" ]] || die "ESI-Bench requires Linux x86_64." + require_commit "$esi_root" "$ESI_BENCH_COMMIT" "ESI-Bench" + require_commit "$behavior_root" "$BEHAVIOR_COMMIT" "BEHAVIOR-1K" + py="$(python_bin)" + require_python_version "$py" "3.11" + [[ -n "${CONDA_PREFIX:-}" ]] || die "Activate the target behavior conda environment first." + case "$py" in + "$CONDA_PREFIX"/*) ;; + *) die "PYTHON resolves outside active CONDA_PREFIX=$CONDA_PREFIX: $py" ;; + esac + map_patch="$agl_root/contrib/recipes/shaper/esi_bench/patches/behavior_floor_maps.patch" + if git -C "$behavior_root" apply --check "$map_patch" >/dev/null 2>&1; then + git -C "$behavior_root" apply "$map_patch" + elif ! git -C "$behavior_root" apply --reverse --check "$map_patch" >/dev/null 2>&1; then + die "The official ESI map patch cannot be applied cleanly to $behavior_root" + fi + if [[ "$option" == "--install-behavior" ]]; then + [[ "${SHAPER_ACCEPT_NVIDIA_EULA:-}" == "YES" ]] || \ + die "Set SHAPER_ACCEPT_NVIDIA_EULA=YES after reviewing the NVIDIA Isaac Sim EULA." + [[ "${SHAPER_ACCEPT_BEHAVIOR_DATASET_TOS:-}" == "YES" ]] || \ + die "Set SHAPER_ACCEPT_BEHAVIOR_DATASET_TOS=YES after reviewing the BEHAVIOR dataset terms." + [[ -n "${OMNIGIBSON_DATA_PATH:-}" ]] || die "Set OMNIGIBSON_DATA_PATH to a large data disk." + mkdir -p "$OMNIGIBSON_DATA_PATH" + if [[ ! -x "$CONDA_PREFIX/bin/pip" ]]; then + "$py" -m ensurepip --upgrade + fi + ( + cd "$behavior_root" + export PATH="$CONDA_PREFIX/bin:$PATH" + [[ "$(command -v python)" == "$CONDA_PREFIX/bin/python" ]] || \ + die "BEHAVIOR setup would use Python outside the active conda environment." + [[ "$(command -v pip)" == "$CONDA_PREFIX/bin/pip" ]] || \ + die "BEHAVIOR setup would use pip outside the active conda environment." + bash setup.sh \ + --omnigibson --bddl --dataset \ + --accept-conda-tos --accept-nvidia-eula --accept-dataset-tos \ + --confirm-no-conda --cuda-version "${SHAPER_BEHAVIOR_CUDA_VERSION:-12.8}" + ) + fi + "$py" -m pip install "google-genai==$GOOGLE_GENAI_VERSION" + PYTHONPATH="$agl_root${PYTHONPATH:+:$PYTHONPATH}" \ + ESI_BEHAVIOR_ROOT="$behavior_root" ESI_BENCH_ROOT="$esi_root" \ + ESI_OMNIGIBSON_DATA_ROOT="${OMNIGIBSON_DATA_PATH:-${ESI_OMNIGIBSON_DATA_ROOT:-}}" \ + "$py" - <<'PY' +from pathlib import Path +import os + +from contrib.recipes.shaper.esi_bench.check_env import ( + check_map_generation_patch, + check_omnigibson_assets, + check_runtime_modules, +) +from contrib.recipes.shaper.esi_bench.contracts import check_behavior_source, check_omnigibson_install + +behavior = Path(os.environ["ESI_BEHAVIOR_ROOT"]) +data = Path(os.environ.get("ESI_OMNIGIBSON_DATA_ROOT", "")) +errors = [ + *check_behavior_source(behavior), + *check_omnigibson_install(behavior), + *check_runtime_modules(), + *check_map_generation_patch( + behavior / "asset_pipeline" / "b1k_pipeline" / "usd_conversion" / "make_maps.py" + ), + *check_omnigibson_assets(data), +] +if errors: + raise SystemExit("\n".join(errors)) +print("ESI-Bench behavior environment, modules, map setting, and assets passed.") +PY + ;; + + *) + usage + exit 2 + ;; +esac diff --git a/contrib/recipes/shaper/scripts/checkout_shaper_benchmarks.sh b/contrib/recipes/shaper/scripts/checkout_shaper_benchmarks.sh new file mode 100755 index 000000000..db623c426 --- /dev/null +++ b/contrib/recipes/shaper/scripts/checkout_shaper_benchmarks.sh @@ -0,0 +1,62 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft. All rights reserved. + +set -euo pipefail + +VLABENCH_REPOSITORY="${VLABENCH_REPOSITORY:-https://github.com/OpenMOSS/VLABench.git}" +VLABENCH_COMMIT="cf588fe60c0c7282174fe979f5913170cfe69017" +OPENPI_REPOSITORY="${OPENPI_REPOSITORY:-https://github.com/Shiduo-zh/openpi.git}" +OPENPI_COMMIT="4483d1da6332da44115fe530e4e6fdd89bd57b13" +ESI_BENCH_REPOSITORY="${ESI_BENCH_REPOSITORY:-https://github.com/ESI-Bench/ESI-Bench.git}" +ESI_BENCH_COMMIT="3c1756396f32b1a90c1f72356a7fde45f418e179" +BEHAVIOR_REPOSITORY="${BEHAVIOR_REPOSITORY:-https://github.com/StanfordVL/BEHAVIOR-1K.git}" +BEHAVIOR_COMMIT="67ad490856dd465d4606663106f81673fc8bf4e8" + +destination="${1:-$PWD/shaper-benchmarks}" +mkdir -p "$destination" +destination="$(cd "$destination" && pwd)" + +checkout_repository() { + local repository="$1" + local commit="$2" + local target="$3" + local created=0 + + if [[ ! -d "$target/.git" ]]; then + if [[ -e "$target" ]]; then + if [[ ! -d "$target" || -n "$(find "$target" -mindepth 1 -maxdepth 1 -print -quit)" ]]; then + echo "Refusing to overwrite non-empty path: $target" >&2 + return 2 + fi + fi + git clone --filter=blob:none --no-checkout "$repository" "$target" + created=1 + fi + # A fresh --no-checkout clone reports every tracked file as deleted until its + # first checkout. Only protect pre-existing worktrees from mutation. + if [[ "$created" -eq 0 && -n "$(git -C "$target" status --porcelain --untracked-files=all)" ]]; then + echo "Refusing to change dirty benchmark checkout: $target" >&2 + return 2 + fi + git -C "$target" fetch --depth=1 origin "$commit" + git -C "$target" checkout --detach "$commit" + test "$(git -C "$target" rev-parse HEAD)" = "$commit" +} + +vlabench_target="$destination/VLABench" +openpi_target="$destination/OpenPI" +esi_target="$destination/ESI-Bench" +behavior_target="$destination/BEHAVIOR-1K" + +checkout_repository "$VLABENCH_REPOSITORY" "$VLABENCH_COMMIT" "$vlabench_target" +checkout_repository "$OPENPI_REPOSITORY" "$OPENPI_COMMIT" "$openpi_target" + +checkout_repository "$ESI_BENCH_REPOSITORY" "$ESI_BENCH_COMMIT" "$esi_target" +checkout_repository "$BEHAVIOR_REPOSITORY" "$BEHAVIOR_COMMIT" "$behavior_target" + +printf 'VLABENCH_CHECKOUT=%s\n' "$vlabench_target" +printf 'VLABENCH_ROOT=%s\n' "$vlabench_target/VLABench" +printf 'OPENPI_ROOT=%s\n' "$openpi_target" +printf 'ESI_BENCH_ROOT=%s\n' "$esi_target" +printf 'ESI_BEHAVIOR_ROOT=%s\n' "$behavior_target" +printf '%s\n' "Source checkouts are pinned. Simulator assets, checkpoints, and Python environments are not downloaded." diff --git a/contrib/recipes/shaper/scripts/download_shaper_vlabench_actor.py b/contrib/recipes/shaper/scripts/download_shaper_vlabench_actor.py new file mode 100644 index 000000000..b4775ceb6 --- /dev/null +++ b/contrib/recipes/shaper/scripts/download_shaper_vlabench_actor.py @@ -0,0 +1,76 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Download and verify the frozen VLABench actor checkpoint. + +Run this helper from the Agent Lightning repository root in an environment +that has ``huggingface_hub`` installed. ``HF_TOKEN`` and ``HF_ENDPOINT`` are +honored by ``huggingface_hub`` without exposing credentials on the command +line. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import Sequence + +from contrib.recipes.shaper.vlabench.actor_contract import ( + CHECKPOINT_INFERENCE_PATTERNS, + CHECKPOINT_MANIFEST_SHA256, + CHECKPOINT_REPOSITORY, + CHECKPOINT_REVISION, + checkpoint_manifest_digest, +) + + +def verify_checkpoint(path: Path) -> str: + """Return the verified manifest digest for one local checkpoint.""" + + resolved = path.expanduser().resolve() + digest = checkpoint_manifest_digest(resolved) + if digest != CHECKPOINT_MANIFEST_SHA256: + raise RuntimeError( + f"Checkpoint manifest digest {digest} does not match pinned " f"{CHECKPOINT_MANIFEST_SHA256}." + ) + return digest + + +def main(argv: Sequence[str] | None = None) -> int: + """Download the immutable revision and verify its identity manifests.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("destination", type=Path) + parser.add_argument("--max-workers", type=int, default=8) + parser.add_argument( + "--verify-only", + action="store_true", + help="Do not contact Hugging Face; verify an existing checkpoint directory.", + ) + args = parser.parse_args(argv) + destination = args.destination.expanduser().resolve() + if int(args.max_workers) < 1: + raise ValueError("max-workers must be positive.") + + if not args.verify_only: + try: + from huggingface_hub import snapshot_download + except ImportError as exc: + raise RuntimeError("Install huggingface_hub before downloading the actor checkpoint.") from exc + snapshot_download( + repo_id=CHECKPOINT_REPOSITORY, + revision=CHECKPOINT_REVISION, + local_dir=destination, + max_workers=int(args.max_workers), + allow_patterns=list(CHECKPOINT_INFERENCE_PATTERNS), + ) + + digest = verify_checkpoint(destination) + print(f"checkpoint_repository={CHECKPOINT_REPOSITORY}") + print(f"checkpoint_revision={CHECKPOINT_REVISION}") + print(f"checkpoint_manifest_sha256={digest}") + print(f"checkpoint_path={destination}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/contrib/recipes/shaper/scripts/run_esi_bench.sh b/contrib/recipes/shaper/scripts/run_esi_bench.sh new file mode 100755 index 000000000..39e204108 --- /dev/null +++ b/contrib/recipes/shaper/scripts/run_esi_bench.sh @@ -0,0 +1,103 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft. All rights reserved. + +set -euo pipefail + +script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +recipe_dir="$(cd "$script_dir/.." && pwd)" +repo_root="$(cd "$recipe_dir/../../.." && pwd)" + +# Configuration. Edit these defaults or export the variables before running. +export SHAPER_BENCH_ROOT="${SHAPER_BENCH_ROOT:-$repo_root/shaper-benchmarks}" +export SHAPER_PLANNER_ENDPOINT="${SHAPER_PLANNER_ENDPOINT:-http://127.0.0.1:8001/v1}" +export SHAPER_MODEL="${SHAPER_MODEL:-Qwen/Qwen3.6-27B}" +export SHAPER_API_KEY_ENV="${SHAPER_API_KEY_ENV:-OPENAI_API_KEY}" +export SHAPER_PLANNER_MAX_TOKENS="${SHAPER_PLANNER_MAX_TOKENS:-32768}" +export SHAPER_OPTIMIZER_MAX_TOKENS="${SHAPER_OPTIMIZER_MAX_TOKENS:-65536}" +export SHAPER_PLANNER_TIMEOUT="${SHAPER_PLANNER_TIMEOUT:-300}" +export SHAPER_PLANNER_RETRIES="${SHAPER_PLANNER_RETRIES:-2}" +export SHAPER_PLANNER_TEMPERATURE="${SHAPER_PLANNER_TEMPERATURE:-1.0}" +export SHAPER_PLANNER_TOP_P="${SHAPER_PLANNER_TOP_P:-0.95}" +export SHAPER_PLANNER_PRESENCE_PENALTY="${SHAPER_PLANNER_PRESENCE_PENALTY:-0.0}" +export SHAPER_PLANNER_EXTRA_BODY="${SHAPER_PLANNER_EXTRA_BODY:-{\"top_k\":20,\"min_p\":0.0,\"repetition_penalty\":1.0}}" + +export ESI_BENCH_ROOT="${ESI_BENCH_ROOT:-$SHAPER_BENCH_ROOT/ESI-Bench}" +export ESI_BEHAVIOR_ROOT="${ESI_BEHAVIOR_ROOT:-$SHAPER_BENCH_ROOT/BEHAVIOR-1K}" +export ESI_WORKER_PYTHON="${ESI_WORKER_PYTHON:-$HOME/miniconda3/envs/shaper-esi-worker/bin/python}" +export ESI_OMNIGIBSON_DATA_ROOT="${ESI_OMNIGIBSON_DATA_ROOT:-$HOME/omnigibson-data}" +export OMNIGIBSON_DATA_PATH="${OMNIGIBSON_DATA_PATH:-$ESI_OMNIGIBSON_DATA_ROOT}" +export ESI_QUESTIONS_JSONL="${ESI_QUESTIONS_JSONL:-$ESI_BENCH_ROOT/hf_dataset/data/questions.jsonl}" +export ESI_MAKE_MAPS_PATH="${ESI_MAKE_MAPS_PATH:-$ESI_BEHAVIOR_ROOT/asset_pipeline/b1k_pipeline/usd_conversion/make_maps.py}" +export ESI_TRAIN_SPLIT="${ESI_TRAIN_SPLIT:-$recipe_dir/esi_bench/splits/recipe_train10.txt}" +export ESI_VALIDATION_SPLIT="${ESI_VALIDATION_SPLIT:-$recipe_dir/esi_bench/splits/recipe_validation10.txt}" + +export ESI_MAX_STEPS="${ESI_MAX_STEPS:-30}" +export ESI_MIN_STEPS="${ESI_MIN_STEPS:-3}" +export ESI_CONFIDENCE_THRESHOLD="${ESI_CONFIDENCE_THRESHOLD:-0.85}" +export ESI_MAX_NEW_TOKENS="${ESI_MAX_NEW_TOKENS:-32768}" +export ESI_TEMPERATURE="${ESI_TEMPERATURE:-$SHAPER_PLANNER_TEMPERATURE}" +export ESI_TOP_P="${ESI_TOP_P:-$SHAPER_PLANNER_TOP_P}" +export ESI_ROBOT="${ESI_ROBOT:-R1}" +export ESI_EPISODE_TIMEOUT="${ESI_EPISODE_TIMEOUT:-1800}" +export ESI_ENVIRONMENT_RETRIES="${ESI_ENVIRONMENT_RETRIES:-1}" + +export SHAPER_HARNESS_TIMEOUT="${SHAPER_HARNESS_TIMEOUT:-3}" +export SHAPER_HARNESS_MEMORY_MB="${SHAPER_HARNESS_MEMORY_MB:-768}" +export SHAPER_HARNESS_MAX_OUTPUT_CHARS="${SHAPER_HARNESS_MAX_OUTPUT_CHARS:-24000000}" + +if [[ -n "${PYTHON:-}" ]]; then + python_bin="$PYTHON" +elif command -v python >/dev/null 2>&1; then + python_bin="$(command -v python)" +elif command -v python3 >/dev/null 2>&1; then + python_bin="$(command -v python3)" +else + echo "Set PYTHON to a Python interpreter." >&2 + exit 2 +fi +output_dir="${SHAPER_OUTPUT_DIR:-$repo_root/outputs/shaper/esi_bench}" +export ESI_OUTPUT_ROOT="${ESI_OUTPUT_ROOT:-$output_dir/runner}" +n_runners="${SHAPER_N_RUNNERS:-1}" +cd "$repo_root" +command="${1:-}" +if [[ -n "$command" ]]; then + shift +fi + +usage() { + cat <&2 + exit 2 + ;; +esac diff --git a/contrib/recipes/shaper/scripts/run_vlabench.sh b/contrib/recipes/shaper/scripts/run_vlabench.sh new file mode 100755 index 000000000..1c0c3e74b --- /dev/null +++ b/contrib/recipes/shaper/scripts/run_vlabench.sh @@ -0,0 +1,99 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft. All rights reserved. + +set -euo pipefail + +script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +recipe_dir="$(cd "$script_dir/.." && pwd)" +repo_root="$(cd "$recipe_dir/../../.." && pwd)" + +# Configuration. Edit these defaults or export the variables before running. +export SHAPER_BENCH_ROOT="${SHAPER_BENCH_ROOT:-$repo_root/shaper-benchmarks}" +export SHAPER_PLANNER_ENDPOINT="${SHAPER_PLANNER_ENDPOINT:-http://127.0.0.1:8001/v1}" +export SHAPER_MODEL="${SHAPER_MODEL:-Qwen/Qwen3.6-27B}" +export SHAPER_API_KEY_ENV="${SHAPER_API_KEY_ENV:-OPENAI_API_KEY}" +export SHAPER_PLANNER_MAX_TOKENS="${SHAPER_PLANNER_MAX_TOKENS:-32768}" +export SHAPER_OPTIMIZER_MAX_TOKENS="${SHAPER_OPTIMIZER_MAX_TOKENS:-65536}" +export SHAPER_PLANNER_TIMEOUT="${SHAPER_PLANNER_TIMEOUT:-300}" +export SHAPER_PLANNER_RETRIES="${SHAPER_PLANNER_RETRIES:-2}" +export SHAPER_PLANNER_TEMPERATURE="${SHAPER_PLANNER_TEMPERATURE:-1.0}" +export SHAPER_PLANNER_TOP_P="${SHAPER_PLANNER_TOP_P:-0.95}" +export SHAPER_PLANNER_PRESENCE_PENALTY="${SHAPER_PLANNER_PRESENCE_PENALTY:-0.0}" +export SHAPER_PLANNER_EXTRA_BODY="${SHAPER_PLANNER_EXTRA_BODY:-{\"top_k\":20,\"min_p\":0.0,\"repetition_penalty\":1.0}}" + +export VLABENCH_ROOT="${VLABENCH_ROOT:-$SHAPER_BENCH_ROOT/VLABench/VLABench}" +export VLABENCH_TRACK="${VLABENCH_TRACK:-track_4_semantic_instruction}" +export MUJOCO_GL="${MUJOCO_GL:-egl}" +export VLABENCH_VLA_HOST="${VLABENCH_VLA_HOST:-127.0.0.1}" +export VLABENCH_VLA_PORT="${VLABENCH_VLA_PORT:-8000}" +export VLABENCH_ACTOR_ID="${VLABENCH_ACTOR_ID:-vlabench-base}" +export VLABENCH_OPENPI_POLICY_CONFIG="${VLABENCH_OPENPI_POLICY_CONFIG:-pi0_ft_vlabench_primitive}" +export VLABENCH_OBSERVATION_SCHEMA="${VLABENCH_OBSERVATION_SCHEMA:-reported_three_camera}" +export VLABENCH_VLA_REPLAN_STEPS="${VLABENCH_VLA_REPLAN_STEPS:-5}" +export VLABENCH_VLA_TIMEOUT="${VLABENCH_VLA_TIMEOUT:-300}" +export VLABENCH_MAX_STEPS="${VLABENCH_MAX_STEPS:-400}" +export VLABENCH_MAX_VLM_ROUNDS="${VLABENCH_MAX_VLM_ROUNDS:-10}" +export VLABENCH_DEFAULT_ROUND_STEPS="${VLABENCH_DEFAULT_ROUND_STEPS:-200}" +export VLABENCH_MIN_ROUND_STEPS="${VLABENCH_MIN_ROUND_STEPS:-1}" +export VLABENCH_MAX_SUBSTEPS="${VLABENCH_MAX_SUBSTEPS:-1}" +export VLABENCH_RESET_WAIT_STEPS="${VLABENCH_RESET_WAIT_STEPS:-10}" + +export SHAPER_HARNESS_TIMEOUT="${SHAPER_HARNESS_TIMEOUT:-3}" +export SHAPER_HARNESS_MEMORY_MB="${SHAPER_HARNESS_MEMORY_MB:-768}" +export SHAPER_HARNESS_MAX_OUTPUT_CHARS="${SHAPER_HARNESS_MAX_OUTPUT_CHARS:-32000000}" + +if [[ -n "${PYTHON:-}" ]]; then + python_bin="$PYTHON" +elif command -v python >/dev/null 2>&1; then + python_bin="$(command -v python)" +elif command -v python3 >/dev/null 2>&1; then + python_bin="$(command -v python3)" +else + echo "Set PYTHON to a Python interpreter." >&2 + exit 2 +fi +output_dir="${SHAPER_OUTPUT_DIR:-$repo_root/outputs/shaper/vlabench}" +n_runners="${SHAPER_N_RUNNERS:-1}" +cd "$repo_root" +command="${1:-}" +if [[ -n "$command" ]]; then + shift +fi + +usage() { + cat <&2 + exit 2 + ;; +esac diff --git a/contrib/recipes/shaper/scripts/start_shaper_planner_vllm.sh b/contrib/recipes/shaper/scripts/start_shaper_planner_vllm.sh new file mode 100755 index 000000000..637534b05 --- /dev/null +++ b/contrib/recipes/shaper/scripts/start_shaper_planner_vllm.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft. All rights reserved. + +set -euo pipefail + +command -v vllm >/dev/null 2>&1 || { + echo "vllm is not installed in the active environment." >&2 + exit 2 +} + +model="${SHAPER_MODEL:-Qwen/Qwen3.6-27B}" +host="${SHAPER_VLLM_HOST:-0.0.0.0}" +port="${SHAPER_VLLM_PORT:-8001}" +tp_size="${SHAPER_VLLM_TP_SIZE:-8}" +max_model_len="${SHAPER_VLLM_MAX_MODEL_LEN:-262144}" +reasoning_parser="${SHAPER_VLLM_REASONING_PARSER:-qwen3}" + +args=( + serve "$model" + --host "$host" + --port "$port" + --served-model-name "$model" + --tensor-parallel-size "$tp_size" + --max-model-len "$max_model_len" +) +if [[ -n "$reasoning_parser" ]]; then + args+=(--reasoning-parser "$reasoning_parser") +fi + +exec vllm "${args[@]}" "$@" diff --git a/contrib/recipes/shaper/scripts/start_shaper_vlabench_actor.sh b/contrib/recipes/shaper/scripts/start_shaper_vlabench_actor.sh new file mode 100755 index 000000000..2e46a751c --- /dev/null +++ b/contrib/recipes/shaper/scripts/start_shaper_vlabench_actor.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft. All rights reserved. + +set -euo pipefail + +if [[ "$#" -lt 3 || "$#" -gt 5 ]]; then + echo "Usage: start_shaper_vlabench_actor.sh AGL_ROOT OPENPI_ROOT CHECKPOINT_DIR [PORT] [ACTOR_ID]" >&2 + exit 2 +fi + +agl_root="$(cd "$1" && pwd)" +openpi_root="$(cd "$2" && pwd)" +checkpoint_dir="$3" +port="${4:-8000}" +actor_id="${5:-vlabench-base}" +actor_python="$openpi_root/.venv/bin/python" + +[[ -x "$actor_python" ]] || { + echo "Missing OpenPI environment: $actor_python" >&2 + exit 2 +} + +# The actor may share one GPU with a small OpenAI-compatible planner. Disable +# JAX's default up-front reservation; this changes allocation, not inference. +export XLA_PYTHON_CLIENT_PREALLOCATE="${XLA_PYTHON_CLIENT_PREALLOCATE:-false}" + +PYTHONPATH="$agl_root${PYTHONPATH:+:$PYTHONPATH}" exec "$actor_python" -m \ + contrib.recipes.shaper.vlabench.openpi_server \ + --openpi-root "$openpi_root" \ + --policy-config pi0_ft_vlabench_primitive \ + --policy-dir "$checkpoint_dir" \ + --actor-id "$actor_id" \ + --observation-schema reported_three_camera \ + --port "$port" diff --git a/contrib/recipes/shaper/vlabench/__init__.py b/contrib/recipes/shaper/vlabench/__init__.py new file mode 100644 index 000000000..3eea6698c --- /dev/null +++ b/contrib/recipes/shaper/vlabench/__init__.py @@ -0,0 +1,32 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""SHAPER integration for the official VLABench simulator and a frozen VLA actor.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .agent import VLABenchAgent, VLABenchRuntimeConfig + from .dataset import load_reported_protocol_datasets + + +def __getattr__(name: str) -> Any: + """Load simulator-facing exports only when downstream code requests them.""" + + if name in {"VLABenchAgent", "VLABenchRuntimeConfig"}: + from .agent import VLABenchAgent, VLABenchRuntimeConfig + + return {"VLABenchAgent": VLABenchAgent, "VLABenchRuntimeConfig": VLABenchRuntimeConfig}[name] + if name == "load_reported_protocol_datasets": + from .dataset import load_reported_protocol_datasets + + return load_reported_protocol_datasets + raise AttributeError(name) + + +__all__ = [ + "VLABenchAgent", + "VLABenchRuntimeConfig", + "load_reported_protocol_datasets", +] diff --git a/contrib/recipes/shaper/vlabench/actor_contract.py b/contrib/recipes/shaper/vlabench/actor_contract.py new file mode 100644 index 000000000..025bd89e7 --- /dev/null +++ b/contrib/recipes/shaper/vlabench/actor_contract.py @@ -0,0 +1,57 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Dependency-free identity contract for the frozen VLABench actor.""" + +from __future__ import annotations + +import hashlib +from pathlib import Path + +OPENPI_REPOSITORY = "https://github.com/Shiduo-zh/openpi" +OPENPI_COMMIT = "4483d1da6332da44115fe530e4e6fdd89bd57b13" +POLICY_CONFIG = "pi0_ft_vlabench_primitive" + +CHECKPOINT_REPOSITORY = "VLABench/pi0-primitive-10task" +CHECKPOINT_REVISION = "1ad73753a74d5cd97e67856664350f3f0baa21dc" +CHECKPOINT_MANIFEST_SHA256 = "39ef720bc93c4d3ccdd135ed6f4b803b8e7ae721cb8c068d9115fb55849d6203" +CHECKPOINT_MANIFEST_FILES = ( + "_CHECKPOINT_METADATA", + "params/_METADATA", + "params/_sharding", + "params/manifest.ocdbt", + "assets/vlabench/vlabench_ft_primitive/norm_stats.json", +) +CHECKPOINT_INFERENCE_PATTERNS = ( + "_CHECKPOINT_METADATA", + "params/**", + "assets/**", +) + + +def checkpoint_manifest_digest(root: Path) -> str: + """Hash the released checkpoint's small identity-bearing manifests.""" + + digest = hashlib.sha256() + for relative_path in CHECKPOINT_MANIFEST_FILES: + path = root / relative_path + if not path.is_file(): + raise FileNotFoundError(f"Missing checkpoint identity file: {path}") + data = path.read_bytes() + digest.update(relative_path.encode("utf-8")) + digest.update(b"\0") + digest.update(len(data).to_bytes(8, "big")) + digest.update(data) + return digest.hexdigest() + + +__all__ = [ + "CHECKPOINT_MANIFEST_FILES", + "CHECKPOINT_MANIFEST_SHA256", + "CHECKPOINT_INFERENCE_PATTERNS", + "CHECKPOINT_REPOSITORY", + "CHECKPOINT_REVISION", + "OPENPI_COMMIT", + "OPENPI_REPOSITORY", + "POLICY_CONFIG", + "checkpoint_manifest_digest", +] diff --git a/contrib/recipes/shaper/vlabench/agent.py b/contrib/recipes/shaper/vlabench/agent.py new file mode 100644 index 000000000..b691c0d4f --- /dev/null +++ b/contrib/recipes/shaper/vlabench/agent.py @@ -0,0 +1,700 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Hierarchical VLM-VLA Agent Lightning rollout for VLABench. + +The official VLABench environment and reward remain authoritative. SHAPER's +planner issues one natural-language command per round; a frozen OpenPI policy +executes that command through the benchmark's websocket protocol. +""" + +from __future__ import annotations + +import collections +import importlib +import logging +import os +import re +import threading +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Mapping, Protocol, cast + +from agentlightning.litagent import LitAgent +from agentlightning.types import NamedResources, Rollout +from contrib.agentlightning.contrib.shaper import ( + EpisodeMetadata, + RoundRecord, + emit_episode_metadata, + emit_round_record, +) + +from ..common import ( + completion_text, + ensure_jsonable, + image_data_url, + image_part, + normalize_content, + openai_client, + require_llm, + require_prompt, + strip_thinking, + text_part, +) +from .contracts import make_harness_validator +from .openpi_identity import ( + REPORTED_THREE_CAMERA, + SUPPORTED_OBSERVATION_SCHEMAS, + validate_server_metadata, +) + +logger = logging.getLogger(__name__) +_ENVIRONMENT_CREATE_LOCK = threading.Lock() + + +class ContextBuilder(Protocol): + """Restricted context-harness call interface.""" + + def __call__(self, history: list[dict[str, Any]]) -> Any: ... + + +class ActorInfrastructureError(RuntimeError): + """The frozen OpenPI actor service failed independently of a candidate.""" + + +@dataclass(frozen=True) +class VLABenchRuntimeConfig: + """Runtime configuration for one VLABench runner process.""" + + vlabench_root: Path + planner_resource_name: str = "planner_llm" + skill_resource_name: str = "skill" + harness_resource_name: str = "harness" + vla_host: str = "127.0.0.1" + vla_port: int = 8000 + vla_replan_steps: int = 5 + vla_inference_timeout_seconds: float = 300.0 + max_vlm_rounds: int = 10 + default_round_steps: int = 200 + min_round_steps: int = 1 + planner_max_completion_tokens: int = 32_768 + max_substeps: int = 1 + joint_tolerance: float = 0.01 + reset_wait_steps: int = 10 + harness_timeout_seconds: float = 3.0 + harness_memory_limit_mb: int = 768 + harness_max_output_chars: int = 32_000_000 + observation_schema: str = REPORTED_THREE_CAMERA + expected_actor_id: str = "" + expected_policy_config: str = "" + + def __post_init__(self) -> None: + for name in ( + "vla_port", + "vla_replan_steps", + "max_vlm_rounds", + "default_round_steps", + "min_round_steps", + "planner_max_completion_tokens", + "max_substeps", + ): + if int(getattr(self, name)) < 1: + raise ValueError(f"{name} must be positive.") + if self.joint_tolerance <= 0: + raise ValueError("joint_tolerance must be positive.") + if self.vla_inference_timeout_seconds <= 0: + raise ValueError("vla_inference_timeout_seconds must be positive.") + if self.observation_schema not in SUPPORTED_OBSERVATION_SCHEMAS: + raise ValueError( + "observation_schema must be one of " + + ", ".join(sorted(SUPPORTED_OBSERVATION_SCHEMAS)) + + f"; got {self.observation_schema!r}." + ) + identity_fields = (self.expected_actor_id, self.expected_policy_config) + if any(identity_fields) and not all(identity_fields): + raise ValueError("expected_actor_id and expected_policy_config must be configured together.") + + +class _TimedOpenPIClient: + """Pinned OpenPI websocket protocol with finite connect and receive waits.""" + + def __init__(self, host: str, port: int, timeout_seconds: float) -> None: + websocket_client = cast(Any, importlib.import_module("websockets.sync.client")) + msgpack_numpy = cast(Any, importlib.import_module("openpi_client.msgpack_numpy")) + self._timeout_seconds = timeout_seconds + self._packer = msgpack_numpy.Packer() + self._unpack = msgpack_numpy.unpackb + self._connection = websocket_client.connect( + f"ws://{host}:{port}", + compression=None, + max_size=None, + open_timeout=timeout_seconds, + close_timeout=min(timeout_seconds, 10.0), + ) + metadata = self._connection.recv(timeout=timeout_seconds) + self.server_metadata = self._unpack(metadata) + + def infer(self, observation: Mapping[str, Any]) -> Mapping[str, Any]: + self._connection.send(self._packer.pack(dict(observation))) + response = self._connection.recv(timeout=self._timeout_seconds) + if isinstance(response, str): + raise RuntimeError(f"OpenPI server returned an error: {response}") + value: object = self._unpack(response) + if not isinstance(value, Mapping): + raise TypeError("OpenPI response must be a mapping.") + return cast(Mapping[str, Any], value) + + def close(self) -> None: + self._connection.close() + + +class _OpenPIVLAPolicy: + """Frozen OpenPI websocket actor using VLABench's official observation map.""" + + def __init__( + self, + host: str, + port: int, + replan_steps: int, + observation_schema: str, + inference_timeout_seconds: float, + expected_actor_id: str = "", + expected_policy_config: str = "", + ) -> None: + self._host = host + self._port = port + self._client: Any = None + self._replan_steps = replan_steps + if observation_schema not in SUPPORTED_OBSERVATION_SCHEMAS: + raise ValueError(f"Unsupported VLABench observation schema: {observation_schema!r}.") + self._observation_schema = observation_schema + self._inference_timeout_seconds = inference_timeout_seconds + self._expected_actor_id = expected_actor_id + self._expected_policy_config = expected_policy_config + self._actions: collections.deque[Any] = collections.deque() + self._connect_with_retry() + + def _connect(self) -> None: + self._discard_client() + client = _TimedOpenPIClient( + self._host, + self._port, + self._inference_timeout_seconds, + ) + if self._expected_actor_id: + errors = validate_server_metadata( + client.server_metadata, + expected_actor_id=self._expected_actor_id, + expected_policy_config=self._expected_policy_config, + expected_observation_schema=self._observation_schema, + ) + if errors: + client.close() + raise RuntimeError("OpenPI actor identity mismatch: " + "; ".join(errors)) + self._client = client + + def _discard_client(self) -> None: + client = self._client + self._client = None + close = getattr(client, "close", None) + if callable(close): + try: + close() + except Exception: + logger.debug("Ignoring an error while closing a failed OpenPI connection.", exc_info=True) + + def _connect_with_retry(self, attempts: int = 3) -> None: + last_error: BaseException | None = None + for attempt in range(attempts): + try: + self._connect() + return + except Exception as exc: + last_error = exc + self._discard_client() + if attempt + 1 < attempts: + time.sleep(float(2**attempt)) + raise ActorInfrastructureError( + f"Could not connect to OpenPI at {self._host}:{self._port} after {attempts} attempts." + ) from last_error + + def reset(self) -> None: + self._actions.clear() + + def predict(self, observation: Mapping[str, Any], instruction: str) -> tuple[Any, Any, Any]: + import numpy as np + + vlabench_utils = cast(Any, importlib.import_module("VLABench.utils.utils")) + + if not self._actions: + rgb = observation["rgb"] + ee_state = observation["ee_state"] + position = ee_state[:3] - np.array([0.0, -0.4, 0.78]) + state = np.concatenate( + [ + position, + vlabench_utils.quaternion_to_euler(ee_state[3:7]), + np.asarray(ee_state[-1]).reshape(-1), + ] + ) + if len(rgb) < 4: + raise ValueError(f"VLABench actor requires four RGB views, received {len(rgb)}.") + payload = { + "observation/image": rgb[2], + "observation/second_image": rgb[0], + "observation/wrist_image": rgb[3], + "observation/state": state, + "prompt": instruction, + } + last_error: BaseException | None = None + response: Mapping[str, Any] | None = None + for attempt in range(3): + try: + response = self._client.infer(payload) + break + except Exception as exc: + last_error = exc + self._discard_client() + if attempt < 2: + time.sleep(float(2**attempt)) + self._connect_with_retry() + if response is None: + raise RuntimeError("OpenPI inference failed after three connection attempts.") from last_error + action_chunk = response.get("actions") + if action_chunk is None or len(action_chunk) < self._replan_steps: + raise ActorInfrastructureError( + f"OpenPI returned {0 if action_chunk is None else len(action_chunk)} actions; " + f"at least {self._replan_steps} are required." + ) + self._actions.extend(action_chunk[: self._replan_steps]) + + action = np.asarray(self._actions.popleft()) + target_position = action[:3].copy() + np.array([0.0, -0.4, 0.78]) + target_euler = action[3:6] + gripper = np.ones(2) * 0.04 if float(action[-1]) >= 0.1 else np.zeros(2) + return target_position, target_euler, gripper + + def close(self) -> None: + self._discard_client() + + +def _load_environment(config: VLABenchRuntimeConfig, task: Mapping[str, Any]) -> Any: + os.environ["VLABENCH_ROOT"] = str(config.vlabench_root) + os.environ.setdefault("MUJOCO_GL", "egl") + + importlib.import_module("VLABench.robots") + importlib.import_module("VLABench.tasks") + environments = cast(Any, importlib.import_module("VLABench.envs")) + + # VLABench mutates shared task configuration while constructing an + # environment, and concurrent EGL initialization is not thread-safe. + with _ENVIRONMENT_CREATE_LOCK: + return environments.load_env( + str(task["task_name"]), + episode_config=task["episode_config"], + random_init=False, + reset_wait_step=config.reset_wait_steps, + run_mode="eval", + ) + + +def _observable_images(observation: Mapping[str, Any], observation_schema: str) -> tuple[Any, Any]: + if observation_schema != REPORTED_THREE_CAMERA: + raise ValueError(f"Unsupported VLABench observation schema: {observation_schema!r}.") + rgb = observation["rgb"] + return rgb[2], rgb[3] + + +def _observation_parts( + observation: Mapping[str, Any], + observation_schema: str, +) -> list[dict[str, Any]]: + main, wrist = _observable_images(observation, observation_schema) + return [ + text_part("Main camera (third-person)"), + image_part(image_data_url(main)), + text_part("Wrist camera (gripper)"), + image_part(image_data_url(wrist)), + ] + + +def _parse_plan(raw_text: str, fallback: str, default_steps: int) -> tuple[str, str, int]: + visible = strip_thinking(raw_text) + answer = re.search(r"(?:^|\n)\s*Answer\s*:\s*(.+?)(?=\n\s*Steps\s*:|\Z)", visible, re.I | re.S) + steps = re.search(r"(?:^|\n)\s*Steps\s*:\s*(\d+)", visible, re.I) + command = answer.group(1).strip() if answer else fallback.strip() + command = command.splitlines()[0].strip() or fallback.strip() + estimated_steps = int(steps.group(1)) if steps else default_steps + reasoning = visible[: answer.start()].strip() if answer else visible.strip() + return reasoning, command, estimated_steps + + +def _planner_content( + *, + instruction: str, + current_step: int, + round_index: int, + context: Any, + observation: Mapping[str, Any], + observation_schema: str, +) -> list[dict[str, Any]]: + content = [ + text_part( + "VLABench task instruction:\n" + + instruction + + f"\n\nEnvironment step: {current_step}\nPlanner round: {round_index + 1}" + ), + text_part("Observable execution context built by the harness:"), + *normalize_content(context), + text_part("Current observation:"), + *_observation_parts(observation, observation_schema), + ] + return content + + +def _official_progress(environment: Any) -> tuple[bool, float]: + """Read the benchmark reward only for scoring and termination.""" + + try: + progress = float(environment.get_task_progress()) + return progress >= 1.0, min(1.0, max(0.0, progress)) + except Exception: + pass + + task = getattr(environment, "task", None) + physics = getattr(environment, "physics", None) + conditions = getattr(task, "conditions", None) + if conditions is None: + return False, 0.0 + + try: + completed = bool(conditions.is_met(physics)) + except Exception: + completed = False + if completed: + return True, 1.0 + + def normalized_progress(value: Any) -> float | None: + if isinstance(value, tuple) and value: + value = cast(tuple[Any, ...], value)[0] + if isinstance(value, bool): + return 1.0 if value else 0.0 + if not isinstance(value, (int, float)): + return None + return min(1.0, max(0.0, float(value))) + + # VLABench's composite tasks expose one of three official condition + # layouts. Reading them here preserves partial reward when the benchmark's + # convenience wrapper fails; none of these values enter planner context. + condition_history = getattr(conditions, "condition_has_been_met", None) + if condition_history is not None: + history = list(condition_history) + if history: + return False, sum(bool(item) for item in history) / len(history) + return False, 0.0 + + condition_sets = getattr(conditions, "condition_sets", None) + if condition_sets is not None: + progresses: list[float] = [] + for condition_set in condition_sets: + try: + progress = normalized_progress(condition_set.met_progress(physics)) + except Exception: + try: + progress = normalized_progress(condition_set.is_met(physics)) + except Exception: + progress = None + if progress is not None: + progresses.append(progress) + return False, max(progresses, default=0.0) + + met_progress = getattr(conditions, "met_progress", None) + if callable(met_progress): + try: + progress = normalized_progress(met_progress(physics)) + except Exception: + progress = None + if progress is not None: + return False, progress + return False, 0.0 + + +def _is_simulator_failure(exc: BaseException) -> bool: + text = f"{type(exc).__name__}: {exc}".lower() + markers = ( + "physics state is invalid", + "badqacc", + "mujoco fatal", + "egl", + "framebuffer", + "render context", + "failed to initialize", + ) + return any(marker in text for marker in markers) + + +def _visible_instruction(task: Mapping[str, Any], environment: Any) -> str: + """Return only the benchmark-visible natural-language instruction.""" + + override = task.get("instruction") + if isinstance(override, str) and override.strip(): + return override.strip() + return str(environment.task.get_instruction()).strip() + + +class VLABenchAgent(LitAgent[dict[str, Any]]): + """Run SHAPER artifacts around a frozen VLABench/OpenPI stack.""" + + def __init__( + self, + config: VLABenchRuntimeConfig, + *, + environment_loader: Callable[[VLABenchRuntimeConfig, Mapping[str, Any]], Any] = _load_environment, + policy_factory: Callable[[str, int, int, str, float, str, str], Any] = _OpenPIVLAPolicy, + ) -> None: + super().__init__() + self.config = config + self._environment_loader = environment_loader + self._policy_factory = policy_factory + + def _harness_validator(self) -> Any: + return make_harness_validator( + timeout_seconds=self.config.harness_timeout_seconds, + memory_limit_mb=self.config.harness_memory_limit_mb, + max_output_chars=self.config.harness_max_output_chars, + ) + + def rollout(self, task: dict[str, Any], resources: NamedResources, rollout: Rollout) -> float: + import numpy as np + + planner_resource = require_llm(resources, self.config.planner_resource_name, rollout) + skill = require_prompt(resources, self.config.skill_resource_name) + harness_source = require_prompt(resources, self.config.harness_resource_name) + context_builder: ContextBuilder = self._harness_validator().runtime(harness_source) + planner_client = openai_client(planner_resource) + + environment: Any = None + policy: Any = None + history: list[dict[str, Any]] = [] + runtime_errors: list[str] = [] + total_steps = 0 + completed = False + reward = 0.0 + termination_reason = "step_budget" + environment_invalid = False + failure_stage = "environment_startup" + + try: + environment = self._environment_loader(self.config, task) + # Match VLABench's official evaluator, which calls reset after + # load_env returns the constructed environment. + environment.reset() + observation = environment.get_observation(require_pcd=False) + instruction = _visible_instruction(task, environment) + failure_stage = "actor_startup" + policy = self._policy_factory( + self.config.vla_host, + self.config.vla_port, + self.config.vla_replan_steps, + self.config.observation_schema, + self.config.vla_inference_timeout_seconds, + self.config.expected_actor_id, + self.config.expected_policy_config, + ) + + for round_index in range(self.config.max_vlm_rounds): + if total_steps >= int(task.get("max_steps", 400)): + break + before = _observation_parts(observation, self.config.observation_schema) + failure_stage = "harness" + context = context_builder(history) + content = _planner_content( + instruction=instruction, + current_step=total_steps, + round_index=round_index, + context=context, + observation=observation, + observation_schema=self.config.observation_schema, + ) + sampling = planner_resource.sampling_parameters + request: dict[str, Any] = { + "model": planner_resource.model, + "messages": [ + {"role": "system", "content": skill}, + {"role": "user", "content": content}, + ], + "max_completion_tokens": int( + sampling.get("max_completion_tokens", self.config.planner_max_completion_tokens) + ), + } + if "temperature" in sampling: + request["temperature"] = sampling["temperature"] + if "top_p" in sampling: + request["top_p"] = sampling["top_p"] + if "presence_penalty" in sampling: + request["presence_penalty"] = sampling["presence_penalty"] + if isinstance(sampling.get("extra_body"), dict): + request["extra_body"] = sampling["extra_body"] + failure_stage = "planner" + response = cast(Any, planner_client.chat.completions.create(**request)) + raw_text, _ = completion_text(response) + reasoning, command, requested_steps = _parse_plan( + raw_text, + instruction, + self.config.default_round_steps, + ) + remaining = int(task.get("max_steps", 400)) - total_steps + if round_index + 1 == self.config.max_vlm_rounds: + # The reported evaluator stops asking the planner after its + # final round but keeps executing that subgoal until the + # episode-level step budget is exhausted. + round_budget = remaining + else: + round_budget = min(remaining, max(self.config.min_round_steps, requested_steps)) + policy.reset() + round_errors: list[str] = [] + executed = 0 + + for _ in range(round_budget): + try: + failure_stage = "actor" + target_position, target_euler, gripper = policy.predict( + observation, + command, + ) + vlabench_utils = cast(Any, importlib.import_module("VLABench.utils.utils")) + quaternion = vlabench_utils.euler_to_quaternion(*target_euler) + ik_success, joints = environment.robot.get_qpos_from_ee_pos( + physics=environment.physics, + pos=target_position, + quat=quaternion, + ) + if not ik_success: + round_errors.append("IK solver failed for the frozen actor action.") + termination_reason = "ik_failure" + break + full_action = np.concatenate([joints, gripper]) + done = False + failure_stage = "simulator" + for _ in range(self.config.max_substeps): + timestep = environment.step(full_action) + if timestep.last(): + done = True + break + current = np.asarray(environment.task.robot.get_qpos(environment.physics)).reshape(-1) + if float(np.max(np.abs(current - full_action[:7]))) < self.config.joint_tolerance: + break + observation = environment.get_observation(require_pcd=False) + total_steps += 1 + executed += 1 + completed, reward = _official_progress(environment) + if done: + # The pinned VLABench environment has an infinite + # time limit and its task termination hooks return + # true only for successful task conditions. Mirror + # the official evaluator, which treats + # timestep.last() as the authoritative success + # signal even if a composite task's convenience + # progress helper is stale or incomplete. + completed = True + reward = 1.0 + if completed: + termination_reason = "completed" + break + except Exception as exc: + message = f"{type(exc).__name__}: {exc}" + round_errors.append(message) + if _is_simulator_failure(exc): + environment_invalid = True + termination_reason = "simulator_failure" + elif failure_stage == "actor": + raise ActorInfrastructureError( + "The frozen OpenPI actor failed while producing an action." + ) from exc + else: + raise + break + + after = _observation_parts(observation, self.config.observation_schema) + emit_round_record( + RoundRecord( + round_index=round_index, + task_instruction=instruction, + planner_response=reasoning, + command=command, + observation_before=before, + observation_after=after, + context_payload=context, + execution_steps=executed, + action_result={"ik_success": not any("IK solver" in item for item in round_errors)}, + runtime_errors=round_errors, + ) + ) + history.append( + { + "round_index": round_index, + "task_instruction": instruction, + "planner_response": reasoning, + "command": command, + "execution_steps": executed, + "observation_before": ensure_jsonable(before), + "observation_after": ensure_jsonable(after), + "action_result": {"ik_success": not any("IK solver" in item for item in round_errors)}, + "runtime_errors": list(round_errors), + } + ) + runtime_errors.extend(round_errors) + if completed or round_errors: + break + + if not completed and termination_reason == "step_budget" and total_steps < int(task.get("max_steps", 400)): + termination_reason = "round_budget" + except Exception as exc: + runtime_errors.append(f"{type(exc).__name__}: {exc}") + simulator_failure = _is_simulator_failure(exc) + candidate_failure = failure_stage in {"harness", "planner"} + environment_invalid = failure_stage == "environment_startup" or simulator_failure + if simulator_failure: + termination_reason = "simulator_failure" + elif candidate_failure: + termination_reason = { + "harness": "harness_failure", + "planner": "planner_failure", + }[failure_stage] + elif failure_stage == "environment_startup": + termination_reason = "environment_startup_failure" + else: + logger.exception( + "VLABench infrastructure failure for %s", + task.get("task_id", "unknown"), + ) + raise + reward = 0.0 + logger.exception("VLABench rollout %s failed", task.get("task_id", "unknown")) + finally: + if policy is not None: + try: + policy.close() + except Exception as exc: + runtime_errors.append(f"actor cleanup: {type(exc).__name__}: {exc}") + if environment is not None: + try: + environment.close() + except Exception as exc: + runtime_errors.append(f"cleanup: {type(exc).__name__}: {exc}") + + emit_episode_metadata( + EpisodeMetadata( + environment_invalid=environment_invalid, + termination_reason=termination_reason, + runtime_errors=runtime_errors, + extra={ + "task_id": str(task.get("task_id", "")), + "environment_steps": total_steps, + "openpi_observation_schema": self.config.observation_schema, + "openpi_actor_id": self.config.expected_actor_id, + "openpi_policy_config": self.config.expected_policy_config, + }, + ) + ) + return float(reward) diff --git a/contrib/recipes/shaper/vlabench/check_env.py b/contrib/recipes/shaper/vlabench/check_env.py new file mode 100644 index 000000000..477f9fb8c --- /dev/null +++ b/contrib/recipes/shaper/vlabench/check_env.py @@ -0,0 +1,238 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Check VLABench/OpenPI prerequisites without consuming an API request.""" + +from __future__ import annotations + +import argparse +import importlib.util +import os +import platform +import socket +from pathlib import Path +from typing import Mapping, Sequence, cast + +from ..cli import endpoint_socket +from .contracts import check_upstream_source +from .dataset import ( + TRACK_NAME, + TRAIN_EPISODES, + VALIDATION_EPISODES, + load_reported_protocol_datasets, + load_track, + track_path, +) +from .openpi_identity import ( + REPORTED_THREE_CAMERA, + SUPPORTED_OBSERVATION_SCHEMAS, + read_server_metadata, + validate_server_metadata, +) + + +def _configured_xml_paths( + root: Path, + *, + track_name: str, + specification: Mapping[str, Sequence[int]], +) -> tuple[list[Path], list[str]]: + """Return model XMLs named by a deterministic episode specification.""" + + errors: list[str] = [] + try: + track = load_track(root, track_name) + except (FileNotFoundError, TypeError, ValueError) as exc: + return [], [f"Cannot inspect VLABench task assets: {exc}"] + + relative_paths: set[str] = set() + for task_name, indices in specification.items(): + episodes = track.get(task_name) + if episodes is None: + errors.append(f"VLABench asset preflight task {task_name!r} is absent from {track_name!r}.") + continue + for index in indices: + if index < 0 or index >= len(episodes): + errors.append( + f"VLABench asset preflight episode {task_name}/ep_{index:03d} is outside " + f"the track's {len(episodes)} episodes." + ) + continue + episode = cast(Mapping[str, object], episodes[index]) + task_value = episode.get("task") + task = cast(Mapping[str, object], task_value) if isinstance(task_value, Mapping) else None + components = task.get("components") if task is not None else None + if not isinstance(components, list): + errors.append(f"VLABench episode {task_name}/ep_{index:03d} has no component list.") + continue + for raw_component in cast(list[object], components): + if not isinstance(raw_component, Mapping): + continue + component = cast(Mapping[str, object], raw_component) + xml_path = component.get("xml_path") + if isinstance(xml_path, str) and xml_path.strip(): + relative_paths.add(xml_path.strip()) + + assets_root = (root / "assets").resolve() + paths: list[Path] = [] + for relative in sorted(relative_paths): + path = (assets_root / relative).resolve() + if not path.is_relative_to(assets_root): + errors.append(f"VLABench episode declares an asset outside the asset root: {relative!r}.") + continue + paths.append(path) + return paths, errors + + +def check_vlabench_assets( + root: Path, + *, + track_name: str = TRACK_NAME, +) -> list[str]: + """Check real files required by the fixed 15/24 protocol.""" + + required = [ + root / "assets" / "obj" / "meshes" / "table" / "table.xml", + root / "assets" / "obj" / "assets" / "textures" / "wood0.png", + root / "assets" / "scenes" / "default" / "empty.xml", + root / "assets" / "scenes" / "default" / "studyroom" / "studyroom.xml", + ] + specification: dict[str, tuple[int, ...]] = {} + for split in (TRAIN_EPISODES, VALIDATION_EPISODES): + for task_name, indices in split.items(): + specification[task_name] = (*specification.get(task_name, ()), *indices) + configured, errors = _configured_xml_paths( + root, + track_name=track_name, + specification=specification, + ) + required.extend(configured) + missing = sorted({path for path in required if not path.is_file()}) + errors.extend(f"Missing VLABench asset payload: {path}" for path in missing[:20]) + if len(missing) > 20: + errors.append(f"Missing {len(missing) - 20} additional VLABench asset files.") + return errors + + +def _check_socket(label: str, host: str, port: int) -> str | None: + try: + with socket.create_connection((host, port), timeout=2.0): + pass + except OSError as exc: + return f"{label} {host}:{port} is unreachable: {exc}" + return None + + +def check_environment( + *, + root: Path, + track_name: str = TRACK_NAME, + host: str, + port: int, + require_vla: bool, + expected_actor_id: str | None = None, + expected_policy_config: str | None = None, + expected_observation_schema: str = REPORTED_THREE_CAMERA, + planner_endpoint: str | None = None, + require_planner: bool = False, +) -> list[str]: + """Return prerequisite errors; an empty list means the static checks pass.""" + + errors: list[str] = [] + errors.extend(check_upstream_source(root)) + errors.extend(check_vlabench_assets(root, track_name=track_name)) + if platform.system() != "Linux": + errors.append(f"VLABench rollout requires Linux; found {platform.system()} {platform.machine()}.") + if not track_path(root, track_name).is_file(): + errors.append(f"Missing official track: {track_path(root, track_name)}") + if importlib.util.find_spec("VLABench") is None: + errors.append("Python package VLABench is not importable.") + if importlib.util.find_spec("openpi_client") is None: + errors.append("Python package openpi_client is not importable.") + if require_vla: + if expected_observation_schema not in SUPPORTED_OBSERVATION_SCHEMAS: + errors.append( + "VLABENCH_OBSERVATION_SCHEMA must be one of " + + ", ".join(sorted(SUPPORTED_OBSERVATION_SCHEMAS)) + + f"; got {expected_observation_schema!r}." + ) + if not expected_actor_id: + errors.append("Set VLABENCH_ACTOR_ID to the identity declared by the OpenPI launcher.") + if not expected_policy_config: + errors.append("Set VLABENCH_OPENPI_POLICY_CONFIG to the pinned OpenPI policy config name.") + if expected_actor_id and expected_policy_config: + try: + metadata = read_server_metadata(host, port) + except Exception as exc: + errors.append(f"OpenPI websocket endpoint {host}:{port} failed its metadata handshake: {exc}") + else: + errors.extend( + validate_server_metadata( + metadata, + expected_actor_id=expected_actor_id, + expected_policy_config=expected_policy_config, + expected_observation_schema=expected_observation_schema, + ) + ) + if planner_endpoint: + planner_host, planner_port = endpoint_socket(planner_endpoint) + if planner_host is None or planner_port is None: + errors.append(f"SHAPER planner endpoint is not a valid HTTP(S) URL: {planner_endpoint!r}") + else: + endpoint_error = _check_socket("SHAPER planner endpoint", planner_host, planner_port) + if endpoint_error: + errors.append(endpoint_error) + elif require_planner: + errors.append("Set SHAPER_PLANNER_ENDPOINT to an OpenAI-compatible chat-completions base URL.") + if track_path(root, track_name).is_file(): + try: + train, validation = load_reported_protocol_datasets(root, track_name=track_name) + except (FileNotFoundError, KeyError, TypeError, ValueError) as exc: + errors.append(f"Cannot load the fixed VLABench split: {exc}") + else: + if len(train) != 15 or len(validation) != 24: + errors.append(f"Unexpected split sizes: train={len(train)} validation={len(validation)}") + return errors + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--root", type=Path, default=Path(os.environ.get("VLABENCH_ROOT", "."))) + parser.add_argument("--track", default=os.environ.get("VLABENCH_TRACK", TRACK_NAME)) + parser.add_argument("--vla-host", default=os.environ.get("VLABENCH_VLA_HOST", "127.0.0.1")) + parser.add_argument("--vla-port", type=int, default=int(os.environ.get("VLABENCH_VLA_PORT", "8000"))) + parser.add_argument("--planner-endpoint", default=os.environ.get("SHAPER_PLANNER_ENDPOINT")) + parser.add_argument("--actor-id", default=os.environ.get("VLABENCH_ACTOR_ID")) + parser.add_argument( + "--policy-config", + default=os.environ.get("VLABENCH_OPENPI_POLICY_CONFIG"), + ) + parser.add_argument( + "--observation-schema", + choices=sorted(SUPPORTED_OBSERVATION_SCHEMAS), + default=os.environ.get("VLABENCH_OBSERVATION_SCHEMA", REPORTED_THREE_CAMERA), + ) + parser.add_argument("--skip-vla-connect", action="store_true") + parser.add_argument("--skip-planner-connect", action="store_true") + args = parser.parse_args(argv) + errors = check_environment( + root=args.root.expanduser().resolve(), + track_name=str(args.track), + host=str(args.vla_host), + port=int(args.vla_port), + require_vla=not bool(args.skip_vla_connect), + expected_actor_id=str(args.actor_id) if args.actor_id else None, + expected_policy_config=str(args.policy_config) if args.policy_config else None, + expected_observation_schema=str(args.observation_schema), + planner_endpoint=str(args.planner_endpoint) if args.planner_endpoint else None, + require_planner=not bool(args.skip_planner_connect), + ) + if errors: + for error in errors: + print(f"[missing] {error}") + return 2 + print("VLABench SHAPER prerequisites passed.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/contrib/recipes/shaper/vlabench/contracts.py b/contrib/recipes/shaper/vlabench/contracts.py new file mode 100644 index 000000000..22bec38e8 --- /dev/null +++ b/contrib/recipes/shaper/vlabench/contracts.py @@ -0,0 +1,253 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""VLABench-specific artifact contracts shared by optimization and rollout.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from typing import Any + +from contrib.agentlightning.contrib.shaper import PythonHarnessValidator + +from ..common import check_python_api, git_revision, git_tracked_changes, validate_multimodal_harness_output +from .actor_contract import OPENPI_COMMIT, OPENPI_REPOSITORY, POLICY_CONFIG + +UPSTREAM_REPOSITORY = "https://github.com/OpenMOSS/VLABench" +UPSTREAM_COMMIT = "cf588fe60c0c7282174fe979f5913170cfe69017" + +_OPENPI_OBSERVATION_KEYS = frozenset( + { + "observation/image", + "observation/second_image", + "observation/wrist_image", + } +) + + +HARNESS_CONTRACT = """Define exactly `def build_context(history)`. +history is a JSON list containing only observable VLABench planner/VLA records: +round_index, task_instruction, planner_response, command, execution_steps, +observation_before and observation_after (OpenAI text/image_url parts), +observable action_result, and runtime_errors. Return a JSON-serializable string +or bounded text/image_url list. Do not access files, network, simulator state, +depth, segmentation, poses, object metadata, rewards, ground truth, or task +IDs.""" + + +def check_upstream_source(vlabench_root: Path) -> list[str]: + """Validate the source revision and lightweight API used by this adapter.""" + + errors: list[str] = [] + revision = git_revision(vlabench_root) + if revision is None: + errors.append(f"VLABench source is not inside a readable Git checkout: {vlabench_root}") + elif revision != UPSTREAM_COMMIT: + errors.append(f"VLABench revision {revision} does not match pinned {UPSTREAM_COMMIT}.") + changes = git_tracked_changes(vlabench_root) + if changes: + errors.append("VLABench checkout has tracked modifications: " + ", ".join(changes) + ".") + errors.extend( + check_python_api( + vlabench_root / "envs" / "__init__.py", + functions={ + "load_env": { + "task", + "episode_config", + "random_init", + "reset_wait_step", + } + }, + ) + ) + return errors + + +def check_openpi_source(openpi_root: Path) -> list[str]: + """Validate the separately deployed frozen VLA actor implementation.""" + + errors: list[str] = [] + revision = git_revision(openpi_root) + if revision is None: + errors.append(f"OpenPI source is not inside a readable Git checkout: {openpi_root}") + elif revision != OPENPI_COMMIT: + errors.append(f"OpenPI revision {revision} does not match pinned {OPENPI_COMMIT}.") + changes = git_tracked_changes(openpi_root) + if changes: + errors.append("OpenPI checkout has tracked modifications: " + ", ".join(changes) + ".") + errors.extend( + check_python_api( + openpi_root / "src" / "openpi" / "policies" / "policy_config.py", + functions={ + "create_trained_policy": { + "train_config", + "checkpoint_dir", + "default_prompt", + } + }, + ) + ) + errors.extend( + check_python_api( + openpi_root / "src" / "openpi" / "serving" / "websocket_policy_server.py", + class_methods={ + "WebsocketPolicyServer": { + "__init__": {"policy", "host", "port", "metadata"}, + "serve_forever": set(), + } + }, + ) + ) + transform_path = openpi_root / "src" / "openpi" / "policies" / "vlabench_policy.py" + try: + transform_tree = ast.parse(transform_path.read_text(encoding="utf-8")) + except (OSError, SyntaxError) as exc: + errors.append(f"Cannot parse pinned OpenPI VLABench transform {transform_path}: {exc}") + else: + input_class = next( + (node for node in transform_tree.body if isinstance(node, ast.ClassDef) and node.name == "VLABenchInputs"), + None, + ) + call_method = ( + next( + ( + node + for node in input_class.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "__call__" + ), + None, + ) + if input_class is not None + else None + ) + consumed_keys: set[str] = set() + if call_method is not None: + for node in ast.walk(call_method): + if ( + isinstance(node, ast.Subscript) + and isinstance(node.value, ast.Name) + and node.value.id == "data" + and isinstance(node.slice, ast.Constant) + and isinstance(node.slice.value, str) + ): + consumed_keys.add(node.slice.value) + missing_keys = sorted(_OPENPI_OBSERVATION_KEYS - consumed_keys) + if missing_keys: + errors.append( + "Pinned OpenPI VLABenchInputs.__call__ does not consume required observation keys: " + + ", ".join(missing_keys) + + "." + ) + config_path = openpi_root / "src" / "openpi" / "training" / "config.py" + try: + config_source = config_path.read_text(encoding="utf-8") + except OSError as exc: + errors.append(f"Cannot read pinned OpenPI training config {config_path}: {exc}") + else: + if f'name="{POLICY_CONFIG}"' not in config_source: + errors.append(f"Pinned OpenPI config {POLICY_CONFIG!r} is missing from {config_path}.") + return errors + + +def validate_skill(source: str) -> list[str]: + """Enforce the planner output contract consumed by the VLABench adapter.""" + + errors: list[str] = [] + if not source.strip(): + errors.append("Skill must not be empty.") + if "Answer:" not in source: + errors.append("Skill must require an `Answer:` line.") + if "Steps:" not in source: + errors.append("Skill must require a `Steps:` line.") + if len(source) > 20_000: + errors.append("Skill must remain below 20,000 characters.") + return errors + + +def _image(label: str) -> dict[str, Any]: + return { + "type": "image_url", + "image_url": {"url": f"data:image/png;base64,{label}"}, + } + + +def _one_round_probe() -> tuple[list[dict[str, Any]]]: + return ( + [ + { + "round_index": 0, + "task_instruction": "Take the red mug and place it in the tray.", + "planner_response": "The mug is visible; approach it.", + "command": "Please take the red mug.", + "execution_steps": 48, + "observation_before": [ + {"type": "text", "text": "Main camera (third-person)"}, + _image("MAIN_BEFORE_0"), + {"type": "text", "text": "Wrist camera (gripper)"}, + _image("WRIST_BEFORE_0"), + ], + "observation_after": [ + {"type": "text", "text": "Main camera (third-person)"}, + _image("MAIN_AFTER_0"), + {"type": "text", "text": "Wrist camera (gripper)"}, + _image("WRIST_AFTER_0"), + ], + "action_result": {"ik_success": True}, + "runtime_errors": [], + } + ], + ) + + +def _two_round_probe() -> tuple[list[dict[str, Any]]]: + history = list(_one_round_probe()[0]) + history.append( + { + "round_index": 1, + "task_instruction": "Take the red mug and place it in the tray.", + "planner_response": "Contact is visible, but placement is incomplete.", + "command": "Please put the red mug into the tray.", + "execution_steps": 72, + "observation_before": history[-1]["observation_after"], + "observation_after": [ + {"type": "text", "text": "Main camera (third-person)"}, + _image("MAIN_AFTER_1"), + {"type": "text", "text": "Wrist camera (gripper)"}, + _image("WRIST_AFTER_1"), + ], + "action_result": {"ik_success": False}, + "runtime_errors": ["IK solver failed once before recovery."], + } + ) + return (history,) + + +def make_harness_validator( + *, + timeout_seconds: float = 3.0, + memory_limit_mb: int = 768, + max_output_chars: int = 32_000_000, +) -> PythonHarnessValidator: + """Build the validator used both when admitting and executing artifacts.""" + + return PythonHarnessValidator( + smoke_args=_one_round_probe(), + additional_smoke_args=(([],), _two_round_probe()), + timeout_seconds=timeout_seconds, + memory_limit_mb=memory_limit_mb, + max_output_chars=max_output_chars, + output_validator=validate_multimodal_harness_output, + ) + + +__all__ = [ + "HARNESS_CONTRACT", + "OPENPI_COMMIT", + "OPENPI_REPOSITORY", + "UPSTREAM_COMMIT", + "UPSTREAM_REPOSITORY", + "check_openpi_source", + "check_upstream_source", + "make_harness_validator", + "validate_skill", +] diff --git a/contrib/recipes/shaper/vlabench/dataset.py b/contrib/recipes/shaper/vlabench/dataset.py new file mode 100644 index 000000000..a151ac52c --- /dev/null +++ b/contrib/recipes/shaper/vlabench/dataset.py @@ -0,0 +1,107 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Deterministic VLABench optimization and validation splits used by SHAPER.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence, cast + +TRACK_NAME = "track_4_semantic_instruction" + +TRAIN_EPISODES: Mapping[str, Sequence[int]] = { + "select_fruit": (0, 1, 2), + "select_toy": (0, 1, 2), + "select_book": (0, 1, 2), + "add_condiment": (0, 1, 2), + "select_painting": (0, 1, 2), +} + +VALIDATION_EPISODES: Mapping[str, Sequence[int]] = { + "select_fruit": (4, 5, 6), + "select_toy": (4, 5, 6), + "select_book": (4, 5, 6), + "add_condiment": (4, 5, 6), + "select_painting": (4, 5, 6), + "select_poker": (1, 2, 3), + "select_mahjong": (1, 2, 3), + "insert_flower": (1, 2, 3), +} + + +def track_path(vlabench_root: Path, track_name: str = TRACK_NAME) -> Path: + """Return the official deterministic track JSON path.""" + + return vlabench_root / "configs" / "evaluation" / "tracks" / f"{track_name}.json" + + +def load_track(vlabench_root: Path, track_name: str = TRACK_NAME) -> dict[str, list[dict[str, Any]]]: + """Load one official VLABench evaluation track.""" + + path = track_path(vlabench_root, track_name) + if not path.is_file(): + raise FileNotFoundError( + f"VLABench track not found at {path}. VLABENCH_ROOT must point to the inner VLABench package directory." + ) + raw: object = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(raw, dict): + raise ValueError(f"VLABench track must be a JSON object: {path}") + output: dict[str, list[dict[str, Any]]] = {} + for task_name, values in cast(dict[str, object], raw).items(): + if not isinstance(values, list): + raise ValueError(f"Invalid episode list for VLABench task {task_name!r}.") + episode_values = cast(list[Any], values) + if not all(isinstance(item, dict) for item in episode_values): + raise ValueError(f"Invalid episode list for VLABench task {task_name!r}.") + output[task_name] = cast(list[dict[str, Any]], episode_values) + return output + + +def materialize_split( + track: Mapping[str, Sequence[dict[str, Any]]], + specification: Mapping[str, Sequence[int]], + *, + max_steps: int, +) -> list[dict[str, Any]]: + """Materialize explicit episode indices as JSON-serializable AGL tasks.""" + + tasks: list[dict[str, Any]] = [] + for task_name, indices in specification.items(): + episodes = track.get(task_name) + if episodes is None: + raise KeyError(f"Task {task_name!r} is absent from the selected VLABench track.") + for index in indices: + if index < 0 or index >= len(episodes): + raise IndexError(f"Episode {index} is out of range for {task_name!r} ({len(episodes)} available).") + tasks.append( + { + "task_id": f"{task_name}/ep_{index:03d}", + "task_name": task_name, + "episode_index": index, + "episode_config": episodes[index], + "max_steps": max_steps, + } + ) + return tasks + + +def load_reported_protocol_datasets( + vlabench_root: Path, + *, + track_name: str = TRACK_NAME, + max_steps: int = 400, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + """Load the reported 15-episode optimization and 24-episode validation protocol.""" + + track = load_track(vlabench_root, track_name) + return ( + materialize_split(track, TRAIN_EPISODES, max_steps=max_steps), + materialize_split(track, VALIDATION_EPISODES, max_steps=max_steps), + ) + + +def task_ids(tasks: Iterable[Mapping[str, Any]]) -> list[str]: + """Return task IDs for diagnostics and split tests.""" + + return [str(task["task_id"]) for task in tasks] diff --git a/contrib/recipes/shaper/vlabench/evaluate.py b/contrib/recipes/shaper/vlabench/evaluate.py new file mode 100644 index 000000000..5bc3c7dac --- /dev/null +++ b/contrib/recipes/shaper/vlabench/evaluate.py @@ -0,0 +1,24 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Evaluate a SHAPER artifact pair on VLABench.""" + +from __future__ import annotations + +from typing import Sequence + +from ..cli import cli_arguments, print_preflight_errors, requests_help +from ..evaluate import main as evaluate_main +from .train import FACTORY, preflight_environment + + +def main(argv: Sequence[str] | None = None) -> int: + arguments = cli_arguments(argv) + if requests_help(arguments): + return evaluate_main(["--factory", FACTORY, *arguments]) + if print_preflight_errors(preflight_environment()): + return 2 + return evaluate_main(["--factory", FACTORY, *arguments]) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/contrib/recipes/shaper/vlabench/factory.py b/contrib/recipes/shaper/vlabench/factory.py new file mode 100644 index 000000000..775bb7d06 --- /dev/null +++ b/contrib/recipes/shaper/vlabench/factory.py @@ -0,0 +1,152 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Environment-configured SHAPER bundle for VLABench.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +from agentlightning.types import LLM, PromptTemplate +from contrib.recipes.shaper.reproduce import ReproductionBundle + +from ..common import load_text +from .actor_contract import ( + CHECKPOINT_MANIFEST_SHA256, + CHECKPOINT_REPOSITORY, + CHECKPOINT_REVISION, +) +from .agent import VLABenchAgent, VLABenchRuntimeConfig +from .contracts import ( + HARNESS_CONTRACT, + OPENPI_COMMIT, + OPENPI_REPOSITORY, + UPSTREAM_COMMIT, + UPSTREAM_REPOSITORY, + check_upstream_source, + make_harness_validator, + validate_skill, +) +from .dataset import TRACK_NAME, load_reported_protocol_datasets, task_ids +from .openpi_identity import REPORTED_THREE_CAMERA +from .roles import VLABenchRoleProtocol + +PROMPT_DIR = Path(__file__).parent / "prompts" + + +def _required_env(name: str) -> str: + value = os.environ.get(name, "").strip() + if not value: + raise ValueError(f"Set {name} before building the VLABench SHAPER bundle.") + return value + + +def _planner_resource() -> LLM: + api_key_env = os.environ.get("SHAPER_API_KEY_ENV", "OPENAI_API_KEY") + sampling: dict[str, Any] = { + "max_completion_tokens": int(os.environ.get("SHAPER_PLANNER_MAX_TOKENS", "32768")), + "optimizer_max_completion_tokens": int(os.environ.get("SHAPER_OPTIMIZER_MAX_TOKENS", "65536")), + "timeout": float(os.environ.get("SHAPER_PLANNER_TIMEOUT", "300")), + "max_retries": int(os.environ.get("SHAPER_PLANNER_RETRIES", "2")), + "temperature": float(os.environ.get("SHAPER_PLANNER_TEMPERATURE", "1.0")), + "top_p": float(os.environ.get("SHAPER_PLANNER_TOP_P", "0.95")), + "presence_penalty": float(os.environ.get("SHAPER_PLANNER_PRESENCE_PENALTY", "0.0")), + } + extra_body = os.environ.get("SHAPER_PLANNER_EXTRA_BODY", "").strip() + if extra_body: + parsed: object = json.loads(extra_body) + if not isinstance(parsed, dict): + raise ValueError("SHAPER_PLANNER_EXTRA_BODY must be a JSON object.") + sampling["extra_body"] = parsed + return LLM( + endpoint=_required_env("SHAPER_PLANNER_ENDPOINT"), + model=_required_env("SHAPER_MODEL"), + api_key=os.environ.get(api_key_env), + sampling_parameters=sampling, + ) + + +def build_bundle() -> ReproductionBundle[dict[str, Any]]: + """Build the complete VLABench training bundle from environment variables.""" + + runtime = VLABenchRuntimeConfig( + vlabench_root=Path(_required_env("VLABENCH_ROOT")).expanduser().resolve(), + vla_host=os.environ.get("VLABENCH_VLA_HOST", "127.0.0.1"), + vla_port=int(os.environ.get("VLABENCH_VLA_PORT", "8000")), + vla_replan_steps=int(os.environ.get("VLABENCH_VLA_REPLAN_STEPS", "5")), + vla_inference_timeout_seconds=float(os.environ.get("VLABENCH_VLA_TIMEOUT", "300")), + max_vlm_rounds=int(os.environ.get("VLABENCH_MAX_VLM_ROUNDS", "10")), + default_round_steps=int(os.environ.get("VLABENCH_DEFAULT_ROUND_STEPS", "200")), + min_round_steps=int(os.environ.get("VLABENCH_MIN_ROUND_STEPS", "1")), + planner_max_completion_tokens=int(os.environ.get("SHAPER_PLANNER_MAX_TOKENS", "32768")), + max_substeps=int(os.environ.get("VLABENCH_MAX_SUBSTEPS", "1")), + joint_tolerance=float(os.environ.get("VLABENCH_JOINT_TOLERANCE", "0.01")), + reset_wait_steps=int(os.environ.get("VLABENCH_RESET_WAIT_STEPS", "10")), + harness_timeout_seconds=float(os.environ.get("SHAPER_HARNESS_TIMEOUT", "3")), + harness_memory_limit_mb=int(os.environ.get("SHAPER_HARNESS_MEMORY_MB", "768")), + harness_max_output_chars=int(os.environ.get("SHAPER_HARNESS_MAX_OUTPUT_CHARS", "32000000")), + observation_schema=os.environ.get("VLABENCH_OBSERVATION_SCHEMA", REPORTED_THREE_CAMERA), + expected_actor_id=_required_env("VLABENCH_ACTOR_ID"), + expected_policy_config=_required_env("VLABENCH_OPENPI_POLICY_CONFIG"), + ) + source_errors = check_upstream_source(runtime.vlabench_root) + if source_errors: + raise RuntimeError("Unsupported VLABench checkout: " + "; ".join(source_errors)) + track_name = os.environ.get("VLABENCH_TRACK", TRACK_NAME) + train, validation = load_reported_protocol_datasets( + runtime.vlabench_root, + track_name=track_name, + max_steps=int(os.environ.get("VLABENCH_MAX_STEPS", "400")), + ) + validator = make_harness_validator( + timeout_seconds=runtime.harness_timeout_seconds, + memory_limit_mb=runtime.harness_memory_limit_mb, + max_output_chars=runtime.harness_max_output_chars, + ) + planner = _planner_resource() + resources = { + "planner_llm": planner, + "skill": PromptTemplate(template=load_text(PROMPT_DIR, "seed_skill.txt"), engine="f-string"), + "harness": PromptTemplate(template=load_text(PROMPT_DIR, "seed_harness.py"), engine="f-string"), + } + return ReproductionBundle( + agent=VLABenchAgent(runtime), + train_dataset=train, + val_dataset=validation, + initial_resources=resources, + planner_resource_name="planner_llm", + harness_contract=HARNESS_CONTRACT, + skill_validator=validate_skill, + harness_validator=validator, + role_protocol=VLABenchRoleProtocol(PROMPT_DIR), + provenance={ + "implementation_scope": "SHAPER method implementation with a benchmark-specific interface and prompt pack", + "benchmark": "VLABench", + "upstream_repository": UPSTREAM_REPOSITORY, + "upstream_commit": UPSTREAM_COMMIT, + "openpi_repository": OPENPI_REPOSITORY, + "openpi_commit": OPENPI_COMMIT, + "checkpoint_repository": CHECKPOINT_REPOSITORY, + "checkpoint_revision": CHECKPOINT_REVISION, + "checkpoint_manifest_sha256": CHECKPOINT_MANIFEST_SHA256, + "split_status": "reported 15-episode optimization and 24-episode fixed-validation protocol", + "track": track_name, + "train_task_ids": task_ids(train), + "validation_task_ids": task_ids(validation), + "actor": { + "type": "frozen OpenPI websocket policy", + "replan_steps": runtime.vla_replan_steps, + "inference_timeout_seconds": runtime.vla_inference_timeout_seconds, + "observation_schema": runtime.observation_schema, + "actor_id": runtime.expected_actor_id, + "policy_config": runtime.expected_policy_config, + }, + "reward": { + "optimization": "official VLABench task progress", + "terminal_success": "official timestep.last() or progress >= 1.0", + }, + "prompt_pack": "contrib/recipes/shaper/vlabench/prompts", + }, + ) diff --git a/contrib/recipes/shaper/vlabench/openpi_identity.py b/contrib/recipes/shaper/vlabench/openpi_identity.py new file mode 100644 index 000000000..a5f9cb9ff --- /dev/null +++ b/contrib/recipes/shaper/vlabench/openpi_identity.py @@ -0,0 +1,82 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Identity handshake for the frozen OpenPI actor used by VLABench.""" + +from __future__ import annotations + +import importlib +from typing import Any, Mapping, cast + +from .actor_contract import ( + CHECKPOINT_MANIFEST_SHA256, + CHECKPOINT_REPOSITORY, + CHECKPOINT_REVISION, + OPENPI_COMMIT, +) + +METADATA_KEY = "shaper_actor" +REPORTED_THREE_CAMERA = "reported_three_camera" +SUPPORTED_OBSERVATION_SCHEMAS = frozenset({REPORTED_THREE_CAMERA}) + + +def read_server_metadata(host: str, port: int, timeout_seconds: float = 5.0) -> Mapping[str, Any]: + """Read the first websocket message without issuing an actor inference.""" + + websocket_client = cast(Any, importlib.import_module("websockets.sync.client")) + msgpack_numpy = cast(Any, importlib.import_module("openpi_client.msgpack_numpy")) + connection = websocket_client.connect( + f"ws://{host}:{port}", + compression=None, + max_size=None, + open_timeout=timeout_seconds, + close_timeout=min(timeout_seconds, 5.0), + ) + try: + metadata: object = msgpack_numpy.unpackb(connection.recv(timeout=timeout_seconds)) + finally: + connection.close() + if not isinstance(metadata, Mapping): + raise TypeError("OpenPI server metadata must be a mapping.") + return cast(Mapping[str, Any], metadata) + + +def validate_server_metadata( + metadata: Mapping[str, Any], + *, + expected_actor_id: str, + expected_policy_config: str, + expected_observation_schema: str, +) -> list[str]: + """Validate metadata emitted by the bundled pinned OpenPI launcher.""" + + raw_identity = metadata.get(METADATA_KEY) + if not isinstance(raw_identity, Mapping): + return [ + "OpenPI server does not expose SHAPER actor identity metadata. Start it with " + "contrib.recipes.shaper.vlabench.openpi_server." + ] + identity = cast(Mapping[str, Any], raw_identity) + errors: list[str] = [] + expected = { + "protocol_version": 1, + "actor_id": expected_actor_id, + "openpi_commit": OPENPI_COMMIT, + "checkpoint_repository": CHECKPOINT_REPOSITORY, + "checkpoint_revision": CHECKPOINT_REVISION, + "checkpoint_manifest_sha256": CHECKPOINT_MANIFEST_SHA256, + "policy_config": expected_policy_config, + "observation_schema": expected_observation_schema, + } + for key, value in expected.items(): + if identity.get(key) != value: + errors.append(f"OpenPI actor metadata {key}={identity.get(key)!r}; expected {value!r}.") + return errors + + +__all__ = [ + "METADATA_KEY", + "REPORTED_THREE_CAMERA", + "SUPPORTED_OBSERVATION_SCHEMAS", + "read_server_metadata", + "validate_server_metadata", +] diff --git a/contrib/recipes/shaper/vlabench/openpi_server.py b/contrib/recipes/shaper/vlabench/openpi_server.py new file mode 100644 index 000000000..448a22b45 --- /dev/null +++ b/contrib/recipes/shaper/vlabench/openpi_server.py @@ -0,0 +1,135 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Serve a pinned VLABench OpenPI checkpoint with verifiable metadata. + +Run this module in the pinned OpenPI uv environment, not the VLABench simulator +environment. It deliberately has no Agent Lightning dependency. +""" + +from __future__ import annotations + +import argparse +import importlib +import logging +import socket +import subprocess +import sys +from pathlib import Path +from typing import Any, Sequence, cast + +from .actor_contract import ( + CHECKPOINT_MANIFEST_SHA256, + CHECKPOINT_REPOSITORY, + CHECKPOINT_REVISION, + OPENPI_COMMIT, + POLICY_CONFIG, + checkpoint_manifest_digest, +) + +METADATA_KEY = "shaper_actor" +SUPPORTED_OBSERVATION_SCHEMAS = ("reported_three_camera",) + + +def _revision(root: Path) -> str: + result = subprocess.run( + ["git", "-C", str(root), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + timeout=5, + ) + return result.stdout.strip() + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--openpi-root", type=Path, required=True) + parser.add_argument("--policy-config", required=True) + parser.add_argument("--policy-dir", required=True) + parser.add_argument("--actor-id", required=True) + parser.add_argument( + "--observation-schema", + choices=SUPPORTED_OBSERVATION_SCHEMAS, + default="reported_three_camera", + ) + parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--default-prompt") + args = parser.parse_args(argv) + + root = cast(Path, args.openpi_root).expanduser().resolve() + revision = _revision(root) + if revision != OPENPI_COMMIT: + raise RuntimeError(f"OpenPI revision {revision} does not match pinned {OPENPI_COMMIT}.") + actor_id = str(args.actor_id).strip() + policy_config_name = str(args.policy_config).strip() + policy_dir_argument = str(args.policy_dir).strip() + observation_schema = str(args.observation_schema) + if not actor_id or not policy_config_name or not policy_dir_argument: + raise ValueError("actor-id, policy-config, and policy-dir must be non-empty.") + if policy_config_name != POLICY_CONFIG: + raise ValueError( + "The bundled VLABench actor launcher requires the paper protocol config " f"{POLICY_CONFIG!r}." + ) + policy_root = Path(policy_dir_argument).expanduser().resolve() + manifest_digest = checkpoint_manifest_digest(policy_root) + if manifest_digest != CHECKPOINT_MANIFEST_SHA256: + raise RuntimeError( + "Checkpoint manifest digest " f"{manifest_digest} does not match pinned {CHECKPOINT_MANIFEST_SHA256}." + ) + policy_dir = str(policy_root) + + source_root = root / "src" + if str(source_root) not in sys.path: + sys.path.insert(0, str(source_root)) + openpi_policy_config = cast( + Any, + importlib.import_module("openpi.policies.policy_config"), + ) + websocket_policy_server = cast( + Any, + importlib.import_module("openpi.serving.websocket_policy_server"), + ) + openpi_config = cast( + Any, + importlib.import_module("openpi.training.config"), + ) + + train_config = openpi_config.get_config(policy_config_name) + policy = openpi_policy_config.create_trained_policy( + train_config, + policy_dir, + default_prompt=cast(str | None, args.default_prompt), + ) + metadata: dict[str, Any] = dict(policy.metadata) + metadata[METADATA_KEY] = { + "protocol_version": 1, + "actor_id": actor_id, + "openpi_commit": OPENPI_COMMIT, + "checkpoint_repository": CHECKPOINT_REPOSITORY, + "checkpoint_revision": CHECKPOINT_REVISION, + "checkpoint_manifest_sha256": CHECKPOINT_MANIFEST_SHA256, + "policy_config": policy_config_name, + "observation_schema": observation_schema, + } + hostname = socket.gethostname() + logging.info( + "Serving SHAPER actor %s with %s from %s on %s:%d", + actor_id, + policy_config_name, + policy_dir, + hostname, + int(args.port), + ) + server = websocket_policy_server.WebsocketPolicyServer( + policy=policy, + host="0.0.0.0", + port=int(args.port), + metadata=metadata, + ) + server.serve_forever() + return 0 + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, force=True) + raise SystemExit(main()) diff --git a/contrib/recipes/shaper/vlabench/prompts/episode_summarizer.txt b/contrib/recipes/shaper/vlabench/prompts/episode_summarizer.txt new file mode 100644 index 000000000..3d6d39f34 --- /dev/null +++ b/contrib/recipes/shaper/vlabench/prompts/episode_summarizer.txt @@ -0,0 +1,23 @@ +You are an episode summarizer for robot manipulation tasks. Your job is to analyze a complete episode execution and produce a concise summary that helps an optimizer understand what went well and what went wrong. + +## Input + +You will receive: +- The original task instruction from the environment +- A list of all sub-tasks the planner issued (one per round) +- The judger's per-round critiques (success/partial/failed + analysis) +- Whether the episode succeeded overall + +## Your Job + +Produce a concise summary (~150-250 words) that identifies: + +1. **Outcome**: Did the episode succeed? How many rounds were used? +2. **What the planner did**: Across rounds, what sub-tasks did the planner issue? How did it interpret the environment instruction? +3. **Cross-round patterns**: Did the planner make progress, get stuck, oscillate between targets, or drift away from its initial choice? +4. **Failure mode (if any)**: If the episode failed, what seems to be the root cause based on the evidence? Infer what you can; don't force a cause. +5. **Context effectiveness**: Was the context information provided to the planner helpful, redundant, or harmful? Did it help the planner track progress or avoid repeating mistakes? + +## Output Format + +Write a single paragraph summary. Be specific — cite exact sub-task wordings when relevant. Do not use bullet points or headers. Describe what happened rather than prescribing fixes — leave the fix-finding to the optimizer. diff --git a/contrib/recipes/shaper/vlabench/prompts/episode_summarizer_user.txt b/contrib/recipes/shaper/vlabench/prompts/episode_summarizer_user.txt new file mode 100644 index 000000000..47d59522a --- /dev/null +++ b/contrib/recipes/shaper/vlabench/prompts/episode_summarizer_user.txt @@ -0,0 +1,6 @@ +Original task instruction: {overall_task} +Episode result: {episode_result} +Total rounds: {total_rounds} + +## Round-by-round execution: +{rounds_detail} diff --git a/contrib/recipes/shaper/vlabench/prompts/harness_optimizer.txt b/contrib/recipes/shaper/vlabench/prompts/harness_optimizer.txt new file mode 100644 index 000000000..5fd8da1fa --- /dev/null +++ b/contrib/recipes/shaper/vlabench/prompts/harness_optimizer.txt @@ -0,0 +1,185 @@ +You are a context code optimizer for the APCO v7 robot manipulation system. + +**IMPORTANT: The planner prompt is FROZEN. You must NOT output a new prompt. Only output improved context code.** + +## System Architecture + +1. **VLM Planner** (Qwen family): Receives task instruction + **multimodal context from build_context()** + current camera views, outputs a sub-task +2. **VLA Executor** (Pi0/VLABench): Executes the sub-task via robot actions +3. **build_context(history)**: Your code. Called every round to build context for the planner. + +## v7: Multimodal Context + +Your `build_context(history)` can return **either**: +- A **string** (plain text context, backward compatible) +- A **list of dicts** (multimodal content blocks with text + images) + +### Multimodal Return Format +```python +[ + {"type": "text", "text": "Previous rounds:"}, + {"type": "text", "text": "--- Round 1 ---\nSub-task: Pick up the banana\nSteps: 50"}, + {"type": "text", "text": "Scene after Round 1:"}, + {"type": "image_url", "image_url": {"url": encode_image(record.obs_after_main)}}, + ... +] +``` + +### Available Helper: `encode_image()` +`encode_image(numpy_array)` → converts a numpy image (H,W,3 uint8) to a `data:image/png;base64,...` URL string. Pre-loaded in global scope. + +## What build_context() Receives + +`history` is a list of `RoundRecord` objects. Each has: +- `round_idx` (int): Zero-based round index +- `vlm_reasoning` (str): The planner's reasoning (after stripping think tags) +- `vlm_raw_output` (str): The planner's full raw output including `` content — contains deeper reasoning, eliminated options, and internal deliberation +- `subtask` (str): The sub-task instruction generated +- `vla_actions_count` (int): Number of VLA steps executed +- `vla_steps_range` (tuple): (start_step, end_step) in the episode +- `obs_before_main` (numpy array or None): Third-person view BEFORE execution +- `obs_after_main` (numpy array or None): Third-person view AFTER execution +- `obs_before_wrist` (numpy array or None): Wrist camera BEFORE execution +- `obs_after_wrist` (numpy array or None): Wrist camera AFTER execution +- `obs_before_state` / `obs_after_state` (numpy array or None): Robot proprioception + +## Token Budget + +Each image costs ~1,000 tokens. The model's context window is large enough to use images, but not free. Choose information richness over premature compression, while still pruning unhelpful repetition. A good default is: keep the most recent 1-3 after-images, summarize older rounds, and include explicit failure/token-drift flags. + +## LLM Calling Capability + +Your context code can freely call an LLM from inside `build_context()`. Use it for **anything** that helps the planner — be creative. The bullets below are illustrative examples, **not an exhaustive list**: +- **Summarization**: compress long text history into a concise state digest +- **Image captioning / visual state extraction**: describe what's in an `obs_after_main` or `obs_after_wrist` image — object positions, gripper contents, container open/close state — as structured text, so the planner gets derived features rather than having to re-parse pixels every round +- **Belief tracking**: extract "what has been tried", "what was learned", "current hypothesis" from `vlm_raw_output` across rounds +- **Failure diagnosis**: given a series of repeated failed subtasks and their resulting images, ask the LLM to propose the likely root cause + +Other uses are welcome: re-ranking past observations by relevance, generating hints for the planner, computing hypotheses the planner should verify, spatial reasoning over images, detecting contradictions between history and current observation — whatever helps. If you can think of a way an LLM call would improve the next planning decision, try it. + +There is **no budget concern** about calling `llm_client` during training; feel free to use it when it creates better state summaries or failure diagnoses. Latency still matters, so keep per-call work bounded. At evaluation time, context code should degrade gracefully if an LLM call fails: return useful text/images rather than crashing. + +Available: +- `llm_client`: Pre-configured OpenAI-compatible client +- `llm_model`: Model name string + +Minimal example — summarising earlier rounds into a short digest: +```python +earlier = history[:-2] # keep last 2 rounds verbatim; summarise the rest +if earlier: + msg = "Summarise these robot manipulation rounds into 3–5 bullet points about progress, stalls, and what's been tried:\n\n" + for rec in earlier: + msg += f"Round {rec.round_idx+1}: {rec.subtask} (executed {rec.vla_actions_count} steps)\n" + resp = llm_client.chat.completions.create( + model=llm_model, + messages=[{"role": "user", "content": msg}], + max_completion_tokens=8192, + temperature=1.0, + ) + digest = resp.choices[0].message.content or "" + parts.append({"type": "text", "text": f"Earlier rounds digest:\n{digest}"}) +``` + +Image captioning example — turn a scene image into structured text the planner can use without re-reading pixels: +```python +img_url = encode_image(rec.obs_after_main) # data URL +resp = llm_client.chat.completions.create( + model=llm_model, + messages=[{"role": "user", "content": [ + {"type": "text", "text": "Describe this scene in <= 60 words. Focus on: gripper contents, object positions, task-relevant state changes."}, + {"type": "image_url", "image_url": {"url": img_url}}, + ]}], + max_completion_tokens=8192, + temperature=1.0, +) +caption = resp.choices[0].message.content or "" +parts.append({"type": "text", "text": f"Scene after Round {rec.round_idx+1}: {caption}"}) +``` + +## Sandbox Contract — strictly enforced + +Your code runs in a restricted Python sandbox. Any violation makes the +entire candidate score 0 and wastes a beam slot. Stay inside this contract. + +### Allowed +- Builtins: `abs all any bin bool bytes callable chr dict dir divmod + enumerate filter float format frozenset hasattr hash hex id int isinstance + issubclass iter len list map max min next object oct ord pow range repr + reversed round set slice sorted str sum tuple type zip True False None` + plus exception types `ValueError TypeError KeyError IndexError + AttributeError StopIteration RuntimeError Exception`. +- Pre-loaded modules (already in scope, do NOT import): `re`, `json`, + `collections`, `math`, `textwrap`, `itertools`, `functools`, `np` (numpy). +- Pre-loaded globals: `encode_image`, `llm_client`, `llm_model`. +- Function argument: `history` (list of `RoundRecord`; passed in when the + sandbox calls your `build_context`). It is NOT a global — declare it via + the function signature only. +- All language features: f-strings, comprehensions, control flow, operators, + user-defined nested functions, `try/except`. + +Note: `hasattr` is allowed but `getattr` is NOT. This asymmetry is intentional +— every `RoundRecord` field is documented above and guaranteed to exist as +an attribute (it may hold `None`), so use direct attribute access. Don't try +to escape via `__builtins__`, `__class__.__mro__`, `__subclasses__()`, or +string-name reflection — the sandbox rejects obvious patterns and silently +corrupts subtle ones. + +### Forbidden — calling any of these makes the candidate fail with score 0 +- Any `import` / `from … import …` statement (the modules above are + already in scope). +- File / process: `open` `exec` `eval` `__import__` `compile` + `.system` `.popen` `.remove` `.rmdir` `.unlink` +- Reflection: `getattr` `setattr` `delattr` `globals` `locals` `vars` +- Misc: `breakpoint` `exit` `quit` `input` `print` + +### Common pitfalls and the safe replacement +- ❌ `getattr(rec, "subtask", "")` ✅ `rec.subtask` — every documented + RoundRecord field is guaranteed to exist (it may be `None`, but the + attribute is always present). Direct attribute access never throws. +- ❌ `getattr(rec, "obs_after_main", None)` ✅ `rec.obs_after_main`, + then check `if rec.obs_after_main is not None:` — `None` is the + documented "no image" value. +- ❌ `hasattr(rec, "x") and getattr(rec, "x")` ✅ `rec.x is not None` +- ❌ `print(...)` for debugging ✅ embed the value into a returned + text block (e.g. `parts.append({"type":"text","text":f"DEBUG: {x}"})`). +- ❌ `import json` / `import re` ✅ already in scope; just use them. +- ❌ `compile(...)` / `eval(...)` for dynamic format strings + ✅ f-strings or `str.format`. +- ❌ Bare `except:` swallowing all errors ✅ catch the specific types + you expect (e.g. `except (KeyError, IndexError):`). Uncaught exceptions + cost this candidate the round, but silent swallowing is worse — the + optimizer thinks your code worked while it actually returned empty + context. +- ❌ `history[-1].subtask` on round 0 ✅ `if history: ... else: ...` — + `history` may be an empty list in round 0; always guard before indexing. + +### Signature & shape +- Define exactly `def build_context(history) -> str | list`. +- `history` is a list of `RoundRecord` with the fields documented above. +- If you return a list, every item must be a dict shaped either + `{"type":"text","text":"..."}` or + `{"type":"image_url","image_url":{"url":"..."}}`. + - ❌ `return ["just a string"]` # rejected: items must be dicts + - ❌ `return [{"text": "..."}]` # rejected: missing "type" + - ✅ `return [{"type": "text", "text": "..."}]` +- Do not hard-code training episode IDs or task-specific success labels. +- Latency matters: keep per-call work bounded; cache LLM calls when + inputs repeat. +- Code that crashes in the sandbox costs a full beam slot. If a + "defensive" pattern is fighting the sandbox, prefer the direct form. + +## Your Task + +Analyze the execution traces and summaries, then produce improved context code that helps the planner make better decisions by providing the right information at the right time. + +## Output Format + +### Analysis + + +### Improved Context Code +```python +def build_context(history): + # Your improved code here + ... +``` diff --git a/contrib/recipes/shaper/vlabench/prompts/harness_optimizer_user.txt b/contrib/recipes/shaper/vlabench/prompts/harness_optimizer_user.txt new file mode 100644 index 000000000..9e760849a --- /dev/null +++ b/contrib/recipes/shaper/vlabench/prompts/harness_optimizer_user.txt @@ -0,0 +1,14 @@ +## Current Planner Prompt (READ-ONLY — do not modify) +``` +{current_prompt} +``` + +## Current Context Management Code +```python +{current_context_code} +``` + +{episode_context} + +## Episode Summaries +{critique} diff --git a/contrib/recipes/shaper/vlabench/prompts/round_judger.txt b/contrib/recipes/shaper/vlabench/prompts/round_judger.txt new file mode 100644 index 000000000..c8f81a224 --- /dev/null +++ b/contrib/recipes/shaper/vlabench/prompts/round_judger.txt @@ -0,0 +1,69 @@ +You are an evaluator for robot manipulation tasks. Analyze task execution results. + +## Input +- 4 images: [BEFORE Main] [BEFORE Wrist] [AFTER Main] [AFTER Wrist] +- Original task instruction from environment +- VLM reasoning and generated sub-task instruction +- Context information that was provided to the planner this round +- Execution statistics (steps taken, progress) + +## Your Role +Objectively evaluate what happened during execution: +1. Did the task succeed, partially succeed, or fail? +2. What changed between BEFORE and AFTER images? +3. What might have caused success or failure? +4. Was the context information useful to the planner, or was it redundant/harmful? + +## Evaluation Criteria + +### Task Completion (5-level stage scoring) +Pick the highest stage clearly supported by the BEFORE/AFTER images and execution stats. Output exactly one of {0.0, 0.25, 0.5, 0.75, 1.0} as `success_score`: +- **0.00 — No progress**: gripper did not move toward the target, or moved in a way unrelated to the goal. +- **0.25 — Reached target**: end-effector is within grasp range of the correct target object (regardless of grasp success). +- **0.50 — Grasped target**: gripper closed on the correct target and is holding it (target lifted or stably gripped). +- **0.75 — Transported toward goal**: target is being held and has been moved into the vicinity of the goal location, but not yet placed within tolerance. +- **1.00 — Completed**: target is placed/oriented within the task's success tolerance, OR the environment reports explicit success. + +Set `subtask_status` consistently: +- `success_score == 1.0` → "success" +- `0.25 <= success_score <= 0.75` → "partial" +- `success_score == 0.0` → "failed" + +The episode-level binary success is decided by the environment, not by you; this stage score is a dense progress signal for the optimizer only. + +### Observation Analysis +Compare BEFORE and AFTER images: +- Object positions: Did target object move? +- Gripper state: Is it holding something? +- Scene changes: Any visible progress toward goal? + +### Reasoning Quality +- Did VLM correctly understand what the task is asking for? +- Did VLM identify the correct target object? +- Is the VLM's reasoning logical? + +### Context Effectiveness +- Did the context provide useful information (e.g., what was already tried, what to do next)? +- Did the planner follow or ignore the context guidance? +- Was the context redundant (repeating what the prompt already says)? +- Was there missing information in the context that could have helped? +- Did the context cause the planner to make worse decisions (e.g., over-constraining, distracting)? + +## Output (JSON only) +```json +{ + "subtask_status": "success" | "partial" | "failed", + "success_score": 0.0 | 0.25 | 0.5 | 0.75 | 1.0, + "observation_analysis": "", + "execution_analysis": "", + "reasoning_analysis": "", + "failure_causes": "", + "improvement_suggestions": "", + "context_analysis": "" +} +``` + +## Important +- Be objective and descriptive +- Focus on observable facts from images +- Note any patterns you observe diff --git a/contrib/recipes/shaper/vlabench/prompts/round_judger_user.txt b/contrib/recipes/shaper/vlabench/prompts/round_judger_user.txt new file mode 100644 index 000000000..4117a28ff --- /dev/null +++ b/contrib/recipes/shaper/vlabench/prompts/round_judger_user.txt @@ -0,0 +1,12 @@ +Task: {overall_task} +Round: {round_idx} +Sub-task: {subtask} +VLA Steps: {vla_steps} + +VLM Reasoning: +{vlm_reasoning} + +Context provided to planner this round: +{context_output} + +[Image 1: BEFORE Main] [Image 2: BEFORE Wrist] [Image 3: AFTER Main] [Image 4: AFTER Wrist] diff --git a/contrib/recipes/shaper/vlabench/prompts/seed_harness.py b/contrib/recipes/shaper/vlabench/prompts/seed_harness.py new file mode 100644 index 000000000..d92d9f68b --- /dev/null +++ b/contrib/recipes/shaper/vlabench/prompts/seed_harness.py @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft. All rights reserved. + + +def build_context(history): + if not history: + return [{"type": "text", "text": "No previous actions (episode start)."}] + parts = [{"type": "text", "text": "Previous rounds of execution:"}] + for record in history: + parts.append( + { + "type": "text", + "text": ( + "Round " + + str(record.get("round_index", 0) + 1) + + ": reasoning=" + + str(record.get("planner_response", "")) + + " | subtask=" + + str(record.get("command", "")) + + " | VLA steps=" + + str(record.get("execution_steps", 0)) + ), + } + ) + observations = record.get("observation_after", []) + for item in observations: + parts.append(item) + return parts diff --git a/contrib/recipes/shaper/vlabench/prompts/seed_skill.txt b/contrib/recipes/shaper/vlabench/prompts/seed_skill.txt new file mode 100644 index 000000000..227d14d27 --- /dev/null +++ b/contrib/recipes/shaper/vlabench/prompts/seed_skill.txt @@ -0,0 +1,21 @@ +You are a robot manipulation planner for a Franka Panda arm. + +## Input +- Two images: [Main camera: third-person view] [Wrist camera: gripper view] +- Task instruction and execution history + +## Environment +VLABench benchmark with primitive manipulation tasks like selecting objects, pressing buttons, and placing items. + +## Your Job +Look at the images and instruction, then output the next sub-task for the robot to execute. + +## Output Format +Brief reasoning about current state and what to do next, then: +Answer: +Steps: + +The "Steps" field indicates how many VLA execution steps this sub-task typically requires: +- Simple movements: 20-50 steps +- Complex or multi-phase actions: 100-200 steps +- Complete tasks: 200-400 steps diff --git a/contrib/recipes/shaper/vlabench/prompts/skill_optimizer.txt b/contrib/recipes/shaper/vlabench/prompts/skill_optimizer.txt new file mode 100644 index 000000000..2d4def966 --- /dev/null +++ b/contrib/recipes/shaper/vlabench/prompts/skill_optimizer.txt @@ -0,0 +1,38 @@ +You are a prompt optimizer for the APCO v7 robot manipulation system on VLABench. + +## System Architecture + +1. **VLM Planner** (Qwen family): Receives task instruction + **multimodal context (text + images from previous rounds)** + current camera views, outputs one VLA sub-task +2. **VLA Executor** (Pi0/VLABench): Executes the sub-task via robot actions +3. **Context Code** (`build_context(history)`): Managed separately. Converts execution history into multimodal context for the planner. You do NOT modify this. + +The planner is called repeatedly during an episode. Each round: planner sees [context history with text/images] + [current observation images] + [task instruction] -> outputs one sub-task -> VLA executes -> repeat. + +## VLA Training Distribution + +The VLA was fine-tuned on VLABench primitive tasks. It understands short, canonical commands that preserve the task's target/container tokens: +- "Put the banana into the plate_seen" +- "Put the mickey into the giftbox_seen" +- "Please select the painting of style ukiyo-e." +- "Add ketchup to the dish" +- "Please take the book contract_law_in_japan" +- "Pick up the 4 of hearts card" +- "Pick up the 1_pin mahjong tile" +- "Place the rose into the vase" +- etc. + +The VLA is weak when the planner invents new object names, rewrites token-like labels, issues verbose strategic instructions, or repeatedly sends the same failed command without adapting. + +## Your Task + +Analyze the execution feedback (episode summaries + planner stats) and produce an improved planner prompt. Focus on systematic failure patterns across episodes. + +Note: The planner now receives multimodal context (text + images from previous rounds). The context code is managed separately — you only optimize the prompt. + +## Output Format + +### Analysis + + +### Improved Prompt + diff --git a/contrib/recipes/shaper/vlabench/prompts/skill_optimizer_user.txt b/contrib/recipes/shaper/vlabench/prompts/skill_optimizer_user.txt new file mode 100644 index 000000000..4f567ab13 --- /dev/null +++ b/contrib/recipes/shaper/vlabench/prompts/skill_optimizer_user.txt @@ -0,0 +1,16 @@ +## Current Prompt +``` +{current_prompt} +``` + +## Current Context Code +```python +{current_context_code} +``` + +The context code above is READ-ONLY for this prompt edit. You may use it to understand what information the planner receives, but your output must be only a new planner prompt. + +{episode_context} + +## Episode Summaries +{critique} diff --git a/contrib/recipes/shaper/vlabench/requirements-simulator.txt b/contrib/recipes/shaper/vlabench/requirements-simulator.txt new file mode 100644 index 000000000..c80f3e6cc --- /dev/null +++ b/contrib/recipes/shaper/vlabench/requirements-simulator.txt @@ -0,0 +1,23 @@ +# Runtime dependencies exercised by the pinned VLABench evaluation adapter. +# The upstream requirements file also includes dataset conversion, notebook, +# and policy-training dependencies that are not imported by simulator rollouts. +setuptools<81 +numpy==1.25.0 +mujoco==3.2.2 +dm-control==1.0.22 +opencv-python-headless==4.10.0.84 +gym==0.26.2 +gymnasium==0.29.1 +mediapy==1.2.0 +open3d==0.18.0 +h5py==3.11.0 +scipy==1.14.0 +scikit-learn==1.5.2 +PyYAML==6.0.2 +networkx==3.3 +colorlog==6.9.0 +colorama==0.4.6 +plotly==5.24.1 +rtree==1.2.0 +gdown==5.2.0 +-e git+https://github.com/motion-planning/rrt-algorithms.git@e51d95ee489a225220d6ae2a764c4111f6ba7d85#egg=rrt-algorithms diff --git a/contrib/recipes/shaper/vlabench/roles.py b/contrib/recipes/shaper/vlabench/roles.py new file mode 100644 index 000000000..a4ffca8fd --- /dev/null +++ b/contrib/recipes/shaper/vlabench/roles.py @@ -0,0 +1,386 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Paper-faithful SHAPER role protocol for VLABench.""" + +from __future__ import annotations + +import asyncio +import json +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Sequence, cast + +from contrib.agentlightning.contrib.shaper import ( + ArtifactProposal, + CandidateEvaluation, + EpisodeTrace, + OptimizationStage, + OptimizerRequestContext, + RoleCompleter, + RoleRequest, + RoundRecord, + parse_json_object, +) + +from ..common import load_text + +_AGL_HARNESS_ADAPTER = """ +## Agent Lightning execution adapter (authoritative for this contrib) + +The paper prompt above describes the original in-process APCO runtime. This +contrib executes generated harnesses in an isolated JSON-only worker. Preserve +the same optimization objective, but obey the following concrete interface: + +- Define exactly `def build_context(history)`. +- `history` is a list of dictionaries. Use `record.get(...)`, not attribute + access. Available keys are `round_index`, `task_instruction`, + `planner_response`, `command`, `execution_steps`, `observation_before`, + `observation_after`, `action_result`, and `runtime_errors`. +- Observations are already OpenAI `text` / `image_url` content parts. Reuse + useful image parts directly; do not call `encode_image`. +- `llm_client`, `llm_model`, `numpy`, and simulator objects are not available. + The harness must be deterministic and may not make network calls. +- The validation contract supplied in the user message overrides incompatible + low-level runtime details in the original prompt. +""".strip() + + +@dataclass(frozen=True) +class _VLABenchGradient: + evaluation: CandidateEvaluation + summaries: tuple[str, ...] + judgements: tuple[tuple[dict[str, Any], ...], ...] + + +def _split_images(value: Any) -> tuple[Any, list[dict[str, Any]]]: + images: list[dict[str, Any]] = [] + + def visit(item: Any) -> Any: + if isinstance(item, list): + return [visit(child) for child in cast(list[Any], item)] + if not isinstance(item, dict): + return item + mapping = cast(dict[str, Any], item) + if mapping.get("type") == "image_url" and isinstance(mapping.get("image_url"), dict): + image = cast(dict[str, Any], mapping["image_url"]) + if isinstance(image.get("url"), str): + images.append(mapping) + return {"type": "image_url", "image_url": {"url": ""}} + return {str(key): visit(child) for key, child in mapping.items()} + + return visit(value), images + + +def _content_parts(value: Any) -> list[dict[str, Any]]: + if isinstance(value, str): + return [{"type": "text", "text": value}] + if not isinstance(value, list): + return [{"type": "text", "text": json.dumps(value, ensure_ascii=False, default=str)}] + output: list[dict[str, Any]] = [] + for item in cast(list[Any], value): + if not isinstance(item, dict): + continue + mapping = cast(dict[str, Any], item) + if mapping.get("type") in {"text", "image_url"}: + output.append(mapping) + return output + + +def _task_instruction(trace: EpisodeTrace) -> str: + return trace.rounds[0].task_instruction if trace.rounds else "" + + +def _episode_success(trace: EpisodeTrace) -> bool: + return float(trace.final_reward or 0.0) >= 1.0 + + +def _clean_markdown_artifact(value: str) -> str: + cleaned = value.strip() + fence = re.fullmatch(r"```(?:text|plaintext|python)?\s*\n?(.*?)```", cleaned, flags=re.DOTALL | re.I) + return fence.group(1).strip() if fence else cleaned + + +class VLABenchRoleProtocol: + """Use the paper's VLABench Judger, Summarizer, and optimizer prompts.""" + + def __init__(self, prompt_dir: Path) -> None: + self.judger_system = load_text(prompt_dir, "round_judger.txt") + self.judger_user = load_text(prompt_dir, "round_judger_user.txt") + self.summarizer_system = load_text(prompt_dir, "episode_summarizer.txt") + self.summarizer_user = load_text(prompt_dir, "episode_summarizer_user.txt") + self.skill_optimizer_system = load_text(prompt_dir, "skill_optimizer.txt") + self.skill_optimizer_user = load_text(prompt_dir, "skill_optimizer_user.txt") + self.harness_optimizer_system = load_text(prompt_dir, "harness_optimizer.txt") + self.harness_optimizer_user = load_text(prompt_dir, "harness_optimizer_user.txt") + + async def _judge_round(self, record: RoundRecord, complete: RoleCompleter) -> dict[str, Any]: + context_text, _ = _split_images(record.context_payload) + user_text = self.judger_user.format( + overall_task=record.task_instruction, + round_idx=record.round_index + 1, + subtask=record.command, + vla_steps=record.execution_steps, + vlm_reasoning=record.planner_response, + context_output=json.dumps(context_text, ensure_ascii=False, default=str), + ) + content: list[dict[str, Any]] = [ + {"type": "text", "text": user_text}, + {"type": "text", "text": "BEFORE observation (Main, then Wrist):"}, + ] + content.extend(record.observation_before) + content.append({"type": "text", "text": "AFTER observation (Main, then Wrist):"}) + content.extend(record.observation_after) + context_parts = _content_parts(record.context_payload) + if context_parts: + content.append( + { + "type": "text", + "text": ( + "Additional multimodal context shown to the planner this round; " + "these are not BEFORE/AFTER transition images:" + ), + } + ) + content.extend(context_parts) + raw = await complete( + RoleRequest( + system_prompt=self.judger_system, + user_content=content, + temperature=1.0, + response_format="json_object", + ) + ) + payload = parse_json_object(raw) + status = str(payload.get("subtask_status", "")).lower() + score = float(payload.get("success_score", -1.0)) + if status not in {"success", "partial", "failed"}: + raise ValueError(f"VLABench Judger returned invalid subtask_status: {status!r}") + if score not in {0.0, 0.25, 0.5, 0.75, 1.0}: + raise ValueError(f"VLABench Judger returned invalid success_score: {score!r}") + required = ( + "observation_analysis", + "execution_analysis", + "reasoning_analysis", + "failure_causes", + "improvement_suggestions", + "context_analysis", + ) + for key in required: + if not isinstance(payload.get(key), str): + raise ValueError(f"VLABench Judger omitted string field {key!r}.") + payload["subtask_status"] = status + payload["success_score"] = score + return payload + + @staticmethod + def _rounds_detail(trace: EpisodeTrace, judgements: Sequence[dict[str, Any]]) -> str: + lines: list[str] = [] + last_index = len(trace.rounds) - 1 + for index, record in enumerate(trace.rounds): + judgement = judgements[index] if index < len(judgements) else {} + lines.append( + f'Round {record.round_index + 1}: subtask="{record.command}" | ' + f"VLA steps={record.execution_steps} | " + f"status={judgement.get('subtask_status', 'unknown')} | " + f"score={judgement.get('success_score', 0.0)}" + ) + lines.append(" Judger: " + json.dumps(judgement, ensure_ascii=False, default=str)) + if index == last_index: + context_text, _ = _split_images(record.context_payload) + lines.append( + " Final context provided to planner: " + json.dumps(context_text, ensure_ascii=False, default=str) + ) + return "\n".join(lines) if lines else "(No planner rounds were emitted.)" + + async def _summarize_episode( + self, + trace: EpisodeTrace, + judgements: Sequence[dict[str, Any]], + complete: RoleCompleter, + ) -> str: + result = "SUCCESS" if _episode_success(trace) else "FAILED" + user_text = self.summarizer_user.format( + overall_task=_task_instruction(trace), + episode_result=f"{result} (progress={float(trace.final_reward or 0.0):.2f})", + total_rounds=len(trace.rounds), + rounds_detail=self._rounds_detail(trace, judgements), + ) + return await complete( + RoleRequest( + system_prompt=self.summarizer_system, + user_content=user_text, + temperature=1.0, + response_format="text", + ) + ) + + async def _diagnose_episode( + self, + trace: EpisodeTrace, + complete: RoleCompleter, + ) -> tuple[str, tuple[dict[str, Any], ...]]: + if trace.rounds: + raw_judgements = await asyncio.gather( + *(self._judge_round(record, complete) for record in trace.rounds), + return_exceptions=True, + ) + judgements: list[dict[str, Any]] = [] + for record, value in zip(trace.rounds, raw_judgements): + if isinstance(value, BaseException): + judgements.append( + { + "subtask_status": "failed", + "success_score": 0.0, + "observation_analysis": "Judger output unavailable.", + "execution_analysis": "Judger output unavailable.", + "reasoning_analysis": "Judger output unavailable.", + "failure_causes": f"Diagnostic failure: {value}", + "improvement_suggestions": "Do not infer a change from this round alone.", + "context_analysis": "Judger output unavailable.", + "round_index": record.round_index, + } + ) + else: + judgements.append(value) + else: + judgements = [] + try: + summary = await self._summarize_episode(trace, judgements, complete) + except (RuntimeError, ValueError) as exc: + summary = ( + f"The episode ended with reward {float(trace.final_reward or 0.0):.2f} after " + f"{len(trace.rounds)} planner rounds. Episode summarization failed: {exc}." + ) + return summary, tuple(judgements) + + async def build_textual_gradient( + self, + evaluation: CandidateEvaluation, + complete: RoleCompleter, + ) -> _VLABenchGradient: + traces = [trace for trace in evaluation.traces if not trace.metadata.environment_invalid] + diagnosed = await asyncio.gather(*(self._diagnose_episode(trace, complete) for trace in traces)) + return _VLABenchGradient( + evaluation=evaluation, + summaries=tuple(item[0] for item in diagnosed), + judgements=tuple(item[1] for item in diagnosed), + ) + + @staticmethod + def _representative_rounds( + evaluation: CandidateEvaluation, limit: int = 4 + ) -> list[tuple[EpisodeTrace, RoundRecord]]: + traces = [trace for trace in evaluation.traces if not trace.metadata.environment_invalid] + selected: list[tuple[EpisodeTrace, RoundRecord]] = [] + first = next(((trace, trace.rounds[0]) for trace in traces if trace.rounds), None) + if first is not None: + selected.append(first) + failures = [(trace, record) for trace in traces if not _episode_success(trace) for record in trace.rounds[1:]] + if failures: + selected.append(max(failures, key=lambda item: item[1].round_index)) + successes = [(trace, record) for trace in traces if _episode_success(trace) for record in trace.rounds[1:]] + if successes: + selected.append(max(successes, key=lambda item: item[1].round_index)) + seen = {(trace.rollout_id, record.round_index) for trace, record in selected} + for trace in traces: + for record in trace.rounds: + key = (trace.rollout_id, record.round_index) + if key not in seen: + selected.append((trace, record)) + seen.add(key) + if len(selected) >= limit: + return selected[:limit] + return selected[:limit] + + @classmethod + def _multimodal_traces(cls, evaluation: CandidateEvaluation) -> list[dict[str, Any]]: + blocks: list[dict[str, Any]] = [ + { + "type": "text", + "text": ( + "## Context Code Execution Traces\n" + "These show what build_context() actually supplied and how the planner responded." + ), + } + ] + for index, (trace, record) in enumerate(cls._representative_rounds(evaluation), start=1): + outcome = "SUCCESS" if _episode_success(trace) else "FAILED" + blocks.append( + { + "type": "text", + "text": ( + f"### Trace {index}: Round {record.round_index + 1}/{len(trace.rounds)}, " + f"episode {outcome}\nbuild_context() output:" + ), + } + ) + blocks.extend(_content_parts(record.context_payload)) + blocks.append({"type": "text", "text": f'Planner subtask: "{record.command}"'}) + return blocks if len(blocks) > 1 else [] + + @staticmethod + def _history(context: OptimizerRequestContext) -> str: + events = [event.model_dump(mode="json") for event in context.optimization_history[-30:]] + return json.dumps(events, ensure_ascii=False, indent=2) + + def build_optimizer_request(self, context: OptimizerRequestContext) -> RoleRequest: + gradient = cast(_VLABenchGradient, context.textual_gradient) + traces = [trace for trace in gradient.evaluation.traces if not trace.metadata.environment_invalid] + successes = sum(_episode_success(trace) for trace in traces) + average_reward = sum(float(trace.final_reward or 0.0) for trace in traces) / len(traces) if traces else 0.0 + episode_context = ( + f"Episodes: {len(traces)}, Success: {successes}/{len(traces)} " + f"({100.0 * successes / max(1, len(traces)):.1f}%), Avg reward: {average_reward:.3f}" + ) + critique = "\n\n".join( + f"Episode {index}: {summary}" for index, summary in enumerate(gradient.summaries, start=1) + ) + template = self.skill_optimizer_user if context.stage == "skill" else self.harness_optimizer_user + user_text = template.format( + current_prompt=context.parent.skill.template, + current_context_code=context.parent.harness.template, + episode_context=episode_context, + critique=critique, + ) + user_text = ( + "## Optimization History (recent)\n" + + self._history(context) + + "\n\n" + + user_text + + f"\n\nProposal round: {context.round_index}; branch: {context.branch_index}." + ) + system_prompt = self.skill_optimizer_system + if context.stage == "harness": + system_prompt = self.harness_optimizer_system + "\n\n" + _AGL_HARNESS_ADAPTER + user_text += ( + "\n\nAGENT LIGHTNING HARNESS VALIDATION CONTRACT\n" + "===========================================\n" + + context.harness_contract + + f"\nFunction: {context.harness_function_name}" + + f"\nSmoke arguments: {list(context.harness_smoke_args)!r}" + ) + if context.validation_feedback: + user_text += context.validation_feedback + content: list[dict[str, Any]] = [{"type": "text", "text": user_text}] + content.extend(self._multimodal_traces(gradient.evaluation)) + return RoleRequest( + system_prompt=system_prompt, + user_content=content, + temperature=1.0, + response_format="text", + ) + + def parse_optimizer_response(self, stage: OptimizationStage, response: str) -> ArtifactProposal: + marker = "### Improved Prompt" if stage == "skill" else "### Improved Context Code" + if marker not in response: + raise ValueError(f"VLABench optimizer response is missing {marker!r}.") + prefix, artifact_section = response.split(marker, maxsplit=1) + analysis = prefix.split("### Analysis", maxsplit=1)[-1].strip() + artifact = _clean_markdown_artifact(artifact_section) + if not artifact: + raise ValueError("VLABench optimizer returned an empty replacement artifact.") + return ArtifactProposal(rationale=analysis or "No analysis supplied.", artifact=artifact) + + +__all__ = ["VLABenchRoleProtocol"] diff --git a/contrib/recipes/shaper/vlabench/train.py b/contrib/recipes/shaper/vlabench/train.py new file mode 100644 index 000000000..f09bf7647 --- /dev/null +++ b/contrib/recipes/shaper/vlabench/train.py @@ -0,0 +1,52 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Preflight and run SHAPER training on the included VLABench adapter.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Sequence + +from ..cli import cli_arguments, print_preflight_errors, requests_help, required_environment +from ..reproduce import main as reproduce_main +from .check_env import check_environment +from .dataset import TRACK_NAME +from .openpi_identity import REPORTED_THREE_CAMERA + +FACTORY = "contrib.recipes.shaper.vlabench.factory:build_bundle" + + +def preflight_environment() -> list[str]: + """Validate the pinned VLABench source, actor identity, and planner.""" + + root = Path(required_environment("VLABENCH_ROOT")).expanduser().resolve() + host = os.environ.get("VLABENCH_VLA_HOST", "127.0.0.1") + port = int(os.environ.get("VLABENCH_VLA_PORT", "8000")) + return check_environment( + root=root, + track_name=os.environ.get("VLABENCH_TRACK", TRACK_NAME), + host=host, + port=port, + require_vla=True, + expected_actor_id=required_environment("VLABENCH_ACTOR_ID"), + expected_policy_config=required_environment("VLABENCH_OPENPI_POLICY_CONFIG"), + expected_observation_schema=os.environ.get("VLABENCH_OBSERVATION_SCHEMA", REPORTED_THREE_CAMERA), + planner_endpoint=required_environment("SHAPER_PLANNER_ENDPOINT"), + require_planner=True, + ) + + +def main(argv: Sequence[str] | None = None) -> int: + arguments = cli_arguments(argv) + if requests_help(arguments): + reproduce_main(["--factory", FACTORY, *arguments]) + return 0 + if print_preflight_errors(preflight_environment()): + return 2 + reproduce_main(["--factory", FACTORY, *arguments]) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())