server/browser_env_environment.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""OpenEnv server environment backed by BrowserGym MiniWoB tasks."""

from __future__ import annotations

import asyncio
import logging
import os
import re
from typing import Any, Dict, List, Mapping, Optional, Tuple
from uuid import uuid4

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

try:
    from ..action_mapping import ActionTranslationError, translate_browser_action
    from ..agents import build_agent_bundle
    from ..artifacts import save_step_artifacts
    from ..curriculum import CurriculumPool, TaskVariant
    from ..models import BrowserAction, BrowserObservation, ConstraintState
    from ..observation import compact_elements_with_stats, extract_instruction, extract_url
    from ..replay import ReplayStore
    from ..reward import RewardEngine
except ImportError:  # pragma: no cover - direct script execution fallback
    from action_mapping import ActionTranslationError, translate_browser_action
    from agents import build_agent_bundle
    from artifacts import save_step_artifacts
    from curriculum import CurriculumPool, TaskVariant
    from models import BrowserAction, BrowserObservation, ConstraintState
    from observation import compact_elements_with_stats, extract_instruction, extract_url
    from replay import ReplayStore
    from reward import RewardEngine


class BrowserGymRuntime:
    """Thin runtime wrapper over real BrowserGym environments."""

    def __init__(self):
        self.env = None
        self._thread_loop: Optional[asyncio.AbstractEventLoop] = None
        self._playwright = None

    def _ensure_thread_event_loop(self) -> None:
        loop: Optional[asyncio.AbstractEventLoop]
        try:
            loop = asyncio.get_event_loop()
        except RuntimeError:
            loop = None
        if loop is None or loop.is_closed():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            self._thread_loop = loop

    def reset(self, variant: TaskVariant) -> Tuple[Mapping[str, Any], Mapping[str, Any]]:
        """Start a fresh BrowserGym environment for one curriculum variant.

        Example:
            runtime = BrowserGymRuntime()
            raw_obs, info = runtime.reset(variant)
        """
        self.close()
        if "miniwob" in variant.task_id and not os.getenv("MINIWOB_URL"):
            raise RuntimeError(
                "MINIWOB_URL is required for BrowserGym MiniWoB tasks. "
                "Install MiniWoB++ and set MINIWOB_URL to the local html/miniwob path."
            )
        self._ensure_thread_event_loop()
        import gymnasium as gym
        import playwright.sync_api
        import browsergym.core
        import browsergym.miniwob  # noqa: F401 - registers MiniWoB envs

        # BrowserGym keeps a process-global Playwright instance. In the OpenEnv server
        # that object can otherwise leak across worker/session threads, which breaks
        # Playwright sync calls with "no running event loop" style errors.
        self._playwright = playwright.sync_api.sync_playwright().start()
        browsergym.core._set_global_playwright(self._playwright)

        headless = os.getenv("BROWSER_ENV_HEADLESS", "1") != "0"
        self.env = gym.make(variant.task_id, headless=headless)
        raw_obs, info = self.env.reset(seed=variant.seed)
        return _ensure_mapping(raw_obs), _ensure_mapping(info)

    def step(self, action: str) -> Tuple[Mapping[str, Any], float, bool, bool, Mapping[str, Any]]:
        """Execute one translated BrowserGym action string."""
        if self.env is None:
            raise RuntimeError("BrowserGym environment must be reset before step")
        raw_obs, reward, terminated, truncated, info = self.env.step(action)
        return _ensure_mapping(raw_obs), float(reward), bool(terminated), bool(truncated), _ensure_mapping(info)

    def close(self) -> None:
        """Tear down the BrowserGym env and the process-global Playwright handle."""
        if self.env is not None:
            try:
                self.env.close()
            except Exception:
                pass
        self.env = None
        if self._playwright is not None:
            try:
                import browsergym.core

                if getattr(browsergym.core, "_PLAYWRIGHT", None) is self._playwright:
                    browsergym.core._set_global_playwright(None)
            except Exception:
                pass
            try:
                self._playwright.stop()
            except Exception:
                pass
            self._playwright = None
        if self._thread_loop is not None:
            try:
                self._thread_loop.close()
            except Exception:
                pass
            try:
                asyncio.set_event_loop(None)
            except Exception:
                pass
            self._thread_loop = None


class BrowserEnvironment(Environment):
    """OpenEnv environment for constraint-aware browser-agent RL.

    This class is the main integration point for the repo:

    - BrowserGym supplies the real browser task runtime
    - the translator resolves structured actions into BrowserGym commands
    - the reward engine applies dense shaping
    - the oracle/judge bundle provides optional LLM assistance and terminal scoring
    - the curriculum/replay layers track progression and artifacts
    """

    SUPPORTS_CONCURRENT_SESSIONS: bool = False

    def __init__(self):
        self.logger = logging.getLogger("browser_env.environment")
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self.curriculum = CurriculumPool()
        self.runtime = BrowserGymRuntime()
        self.reward_engine = RewardEngine()
        self.agents = build_agent_bundle()
        self.llm_services_enabled = os.getenv("BROWSER_ENV_DISABLE_LLM") != "1"
        # Keep curriculum self-evolution deterministic and robust: step-budget
        # mutation works even when oracle/judge LLM services are disabled.
        self.task_generator_enabled = os.getenv("BROWSER_ENV_ENABLE_STEP_BUDGET_ADAPTATION", "1") != "0"
        self.replay = ReplayStore()
        self.current_variant = self.curriculum.current()
        self.raw_obs: Mapping[str, Any] = {}
        self.raw_info: Mapping[str, Any] = {}
        self.history: List[Dict[str, Any]] = []
        self.trajectory: List[Dict[str, Any]] = []
        self.constraints = ConstraintState()
        self.latest_observation_stats: Dict[str, Any] = {}
        self.last_translated_action: Optional[str] = None
        self.consecutive_repeated_actions = 0
        self.consecutive_no_progress_steps = 0
        self.last_progress_evidence: Dict[str, Any] = {}
        self.total_reward = 0.0
        self.episode_done = False
        self.final_success = False
        self.final_failure_reason = "none"
        self.final_failure_detail: Dict[str, Any] = {}
        self.oracle_submit_after_type_repeat = max(
            1,
            int(os.getenv("BROWSER_ENV_ORACLE_SUBMIT_AFTER_TYPE_REPEAT", "1")),
        )
        self.oracle_submit_after_select_repeat = max(
            1,
            int(os.getenv("BROWSER_ENV_ORACLE_SUBMIT_AFTER_SELECT_REPEAT", "1")),
        )
        self.curriculum_advance_on_failure_after = max(
            0,
            int(os.getenv("BROWSER_ENV_CURRICULUM_ADVANCE_ON_FAILURE_AFTER", "3")),
        )
        self.max_repeated_actions = max(
            1,
            int(os.getenv("BROWSER_ENV_MAX_REPEATED_ACTIONS", "3")),
        )
        self.max_invalid_actions = max(
            1,
            int(os.getenv("BROWSER_ENV_MAX_INVALID_ACTIONS", "3")),
        )
        self.max_no_progress_steps_base = max(
            0,
            int(os.getenv("BROWSER_ENV_MAX_NO_PROGRESS_STEPS", "4")),
        )
        self.max_no_progress_steps_cap = max(
            self.max_no_progress_steps_base,
            int(os.getenv("BROWSER_ENV_MAX_NO_PROGRESS_STEPS_CAP", "8") or 8),
        )
        self.no_progress_step_ratio = max(
            0.0,
            float(os.getenv("BROWSER_ENV_NO_PROGRESS_STEP_RATIO", "0.4") or 0.4),
        )
        self.max_no_progress_steps = self.max_no_progress_steps_base
        self.logger.info(
            "env_init episode_id=%s task_id=%s variant=%s llm_enabled=%s",
            self._state.episode_id,
            self.current_variant.task_id,
            self.current_variant.variant_id,
            self.llm_services_enabled,
        )

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        task_id: Optional[str] = None,
        variant_id: Optional[str] = None,
        **_: Any,
    ) -> BrowserObservation:
        """Reset to the next curriculum variant and return the first observation."""
        self._state = State(episode_id=str(uuid4()), step_count=0)
        if episode_id:
            self._state.episode_id = episode_id
        self.current_variant = self.curriculum.variant_for_reset(
            seed,
            task_id=task_id,
            variant_id=variant_id,
        )
        if self.max_no_progress_steps_base <= 0:
            self.max_no_progress_steps = 0
        else:
            scaled = int(round(self.current_variant.max_steps * self.no_progress_step_ratio))
            self.max_no_progress_steps = max(
                self.max_no_progress_steps_base,
                min(self.max_no_progress_steps_cap, scaled),
            )
        self.raw_obs, self.raw_info = self.runtime.reset(self.current_variant)
        self.history = []
        self.trajectory = []
        self.last_translated_action = None
        self.consecutive_repeated_actions = 0
        self.consecutive_no_progress_steps = 0
        self.last_progress_evidence = {}
        self.total_reward = 0.0
        self.episode_done = False
        self.final_success = False
        self.final_failure_reason = "none"
        self.final_failure_detail = {}
        self.constraints = self._new_constraints()
        reset_observation = self._make_observation(
            reward=0.0,
            done=False,
            success=False,
            failure_reason="none",
            metadata={
                "event": "reset",
                "runtime": "browsergym",
                "raw_info": dict(self.raw_info),
                "requested_variant_id": variant_id,
                "requested_task_id": task_id,
            },
        )
        observation_stats = dict(reset_observation.metadata.get("observation_stats") or {})
        self.logger.info(
            "reset episode_id=%s task_id=%s variant=%s difficulty=%s seed=%s max_steps=%s max_no_progress_steps=%s instruction=%r elements=%s",
            self._state.episode_id,
            self.current_variant.task_id,
            self.current_variant.variant_id,
            self.current_variant.difficulty,
            self.current_variant.seed,
            self.current_variant.max_steps,
            self.max_no_progress_steps,
            extract_instruction(self.raw_obs, self.raw_info)[:200],
            observation_stats.get("elements_after", 0),
        )
        return reset_observation

    def step(
        self,
        action: BrowserAction,
        timeout_s: Optional[float] = None,
        **_: Any,
    ) -> BrowserObservation:  # type: ignore[override]
        """Advance the environment by one structured browser action.

        The step flow is intentionally staged:

        1. optionally ask the oracle for help
        2. translate the safe structured action into BrowserGym syntax
        3. execute the BrowserGym step
        4. apply shaped reward components
        5. if terminal, ask the judge to score the full trajectory
        """
        if self._is_done():
            self.logger.warning(
                "step_after_done episode_id=%s step=%s action=%s",
                self._state.episode_id,
                self._state.step_count,
                action.action_type,
            )
            return self._make_observation(
                reward=0.0,
                done=True,
                success=self.final_success,
                failure_reason=self.final_failure_reason,
                metadata={
                    "event": "step_after_done",
                    "failure_detail": dict(self.final_failure_detail),
                },
            )

        pre_step_observation = self._make_observation(
            reward=0.0,
            done=False,
            success=False,
            failure_reason="none",
        )
        action_to_execute = action
        oracle_delta = 0
        current_elements = list(pre_step_observation.elements)
        if action.action_type == "ask_oracle":
            if self.constraints.oracle_calls >= self.constraints.llm_budget:
                self.logger.warning(
                    "oracle_budget_exceeded episode_id=%s step=%s llm_budget=%s",
                    self._state.episode_id,
                    self._state.step_count,
                    self.constraints.llm_budget,
                )
                return self._penalized_invalid_observation(action, "oracle_budget_exceeded")
            oracle_delta = 1
            self.constraints.oracle_calls += 1
            try:
                action_to_execute = self.agents.suggest_action(pre_step_observation, self.history)
            except Exception as exc:
                self.logger.exception(
                    "oracle_failed episode_id=%s step=%s error=%s",
                    self._state.episode_id,
                    self._state.step_count,
                    exc,
                )
                return self._penalized_invalid_observation(action, "browser_error", oracle_delta=oracle_delta)
            action_to_execute = self._repair_oracle_action(
                action_to_execute,
                current_elements,
                instruction=pre_step_observation.instruction,
            )
            self.logger.info(
                "oracle_action episode_id=%s step=%s suggested_action=%s target_id=%s text=%r reasoning=%r",
                self._state.episode_id,
                self._state.step_count,
                action_to_execute.action_type,
                action_to_execute.target_id,
                action_to_execute.text,
                action_to_execute.reasoning,
            )

        invalid_delta = 0
        repetition_delta = 0
        translation_metadata: Dict[str, Any] = {}
        try:
            translation = translate_browser_action(
                action_to_execute,
                current_elements,
                instruction=pre_step_observation.instruction,
                history=self.history,
            )
            translated = translation.browsergym_action
            translation_metadata = dict(translation.metadata)
            action_to_execute = translation.resolved_action
        except ActionTranslationError as exc:
            self.logger.warning(
                "action_translation_error episode_id=%s step=%s action=%s target_id=%s error=%s",
                self._state.episode_id,
                self._state.step_count,
                action_to_execute.action_type,
                action_to_execute.target_id,
                exc,
            )
            return self._penalized_invalid_observation(
                action_to_execute,
                "invalid_action",
                oracle_delta=oracle_delta,
            )

        if translated == self.last_translated_action:
            repetition_delta = 1
            self.constraints.repeated_actions += 1
            self.consecutive_repeated_actions += 1
        else:
            self.consecutive_repeated_actions = 0
        self.last_translated_action = translated

        browser_reward = 0.0
        terminated = False
        truncated = False
        info: Mapping[str, Any] = {}
        failure_reason = "none"
        failure_detail: Dict[str, Any] = {}
        try:
            self.raw_obs, browser_reward, terminated, truncated, info = self.runtime.step(translated)
        except Exception as exc:
            invalid_delta = 1
            self.constraints.invalid_actions += 1
            failure_reason = "browsergym_action_error"
            failure_detail = {
                "error": str(exc),
                "error_type": exc.__class__.__name__,
                "translated_action": translated,
                "target_id": action_to_execute.target_id,
            }
            info = dict(failure_detail)
            self.logger.exception(
                "runtime_step_failed episode_id=%s step=%s action=%s translated=%s target_id=%s error=%s",
                self._state.episode_id,
                self._state.step_count,
                action_to_execute.action_type,
                translated,
                action_to_execute.target_id,
                exc,
            )

        self._state.step_count += 1
        success = bool(terminated and browser_reward > 0)
        submit_like_failure = self._is_submit_like_action(action_to_execute, current_elements) and browser_reward <= 0.0
        progress_evidence = self._assess_progress_evidence(
            previous_elements=current_elements,
            action=action_to_execute,
            browser_reward=browser_reward,
            success=success,
        )
        self.last_progress_evidence = dict(progress_evidence)
        progress_signal = bool(progress_evidence.get("progress_signal"))
        if progress_signal:
            self.consecutive_no_progress_steps = 0
        else:
            self.consecutive_no_progress_steps += 1
        # Treat repetition as terminal only when it is paired with no progress.
        # This avoids aborting useful retries during transient UI/render delays.
        repeated_action_loop = (
            self.consecutive_repeated_actions >= self.max_repeated_actions
            and not progress_signal
        )
        too_many_invalid_actions = self.constraints.invalid_actions >= self.max_invalid_actions
        low_progress_abort = (
            self.max_no_progress_steps > 0
            and self.consecutive_no_progress_steps >= self.max_no_progress_steps
        ) or (
            submit_like_failure
            and not self._has_form_action()
            and not success
        )
        done = bool(
            terminated
            or truncated
            or self._state.step_count >= self.current_variant.max_steps
            or repeated_action_loop
            or too_many_invalid_actions
            or low_progress_abort
        )
        if success:
            failure_reason = "success"
            failure_detail = {"mode": "success"}
        elif failure_reason == "none" and repeated_action_loop:
            failure_reason = "repeated_action_loop"
            failure_detail = {
                "mode": "repeated_action_loop",
                "consecutive_repeated_actions": self.consecutive_repeated_actions,
                "translated_action": translated,
            }
        elif failure_reason == "none" and too_many_invalid_actions:
            failure_reason = "too_many_invalid_actions"
            failure_detail = {
                "mode": "too_many_invalid_actions",
                "invalid_actions": self.constraints.invalid_actions,
            }
        elif failure_reason == "none" and low_progress_abort:
            failure_reason = "low_progress_abort"
            failure_detail = self._build_failure_detail(
                action=action_to_execute,
                translated=translated,
                browser_reward=browser_reward,
                terminated=terminated,
                truncated=truncated,
                submit_like_failure=submit_like_failure,
                mode="low_progress_abort",
            )
        elif failure_reason == "none" and self._state.step_count >= self.current_variant.max_steps:
            failure_reason = "max_steps_exceeded"
            failure_detail = {
                "mode": "max_steps_exceeded",
                "step_count": self._state.step_count,
                "max_steps": self.current_variant.max_steps,
                "translated_action": translated,
            }
        elif failure_reason == "none" and done:
            failure_reason = "submission_failed" if submit_like_failure else "task_failed"
            failure_detail = self._build_failure_detail(
                action=action_to_execute,
                translated=translated,
                browser_reward=browser_reward,
                terminated=terminated,
                truncated=truncated,
                submit_like_failure=submit_like_failure,
                mode=failure_reason,
            )

        delayed_penalty = self._delayed_penalty_for_failure(
            failure_reason=failure_reason,
            failure_detail=failure_detail,
        )
        if delayed_penalty:
            self.constraints.delayed_failures += 1

        reward_breakdown = self.reward_engine.compute(
            browsergym_reward=browser_reward,
            success=success,
            step_delta=1,
            progress_delta=int(progress_evidence.get("reward_credit", 0)),
            oracle_delta=oracle_delta,
            mistake_delta=invalid_delta,
            repetition_delta=repetition_delta,
            delayed_penalty=delayed_penalty,
        )

        event = {
            "step": self._state.step_count,
            "observation": _dump_model(pre_step_observation),
            "action": _dump_action(action),
            "executed_action": _dump_action(action_to_execute),
            "browsergym_action": translated,
            "target_resolution": translation_metadata,
            "browsergym_reward": browser_reward,
            "reward_breakdown": _dump_model(reward_breakdown),
            "success": success,
            "failure_reason": failure_reason,
            "failure_detail": dict(failure_detail),
            "progress_evidence": dict(progress_evidence),
            "constraints": _dump_model(self.constraints),
        }
        self.trajectory.append(event)

        judge_result = None
        if done:
            # The judge is terminal-only today: it sees the full trajectory and
            # adds a final quality reward, while the dense step reward still
            # comes from BrowserGym plus local shaping penalties.
            try:
                judge_result = self.agents.score_trajectory(self.trajectory)
            except Exception as exc:
                self.logger.exception(
                    "judge_failed episode_id=%s step=%s error=%s",
                    self._state.episode_id,
                    self._state.step_count,
                    exc,
                )
                event["judge"] = {
                    "label": "unscored",
                    "reward": 0.0,
                    "rationale": f"judge_failed:{exc.__class__.__name__}",
                }
            else:
                reward_breakdown.trajectory_quality = judge_result.reward
                reward_breakdown.judge_quality_reward = judge_result.reward
                reward_breakdown.total += judge_result.reward
                event["judge"] = {
                    "label": judge_result.label,
                    "reward": judge_result.reward,
                    "rationale": judge_result.rationale,
                }
            event["reward_breakdown"] = _dump_model(reward_breakdown)
        artifact_paths = save_step_artifacts(
            episode_id=self._state.episode_id or "unknown",
            step=self._state.step_count,
            raw_obs=self.raw_obs,
            event=event,
        )
        event["artifacts"] = artifact_paths
        self.history.append(
            {
                "action_type": action_to_execute.action_type,
                "target_id": action_to_execute.target_id,
                "text": action_to_execute.text,
                "browsergym_action": translated,
                "reward": reward_breakdown.total,
                "status": failure_reason,
                "target_resolution": translation_metadata,
                "raw_info": dict(info),
            }
        )
        self.total_reward += reward_breakdown.total
        if done:
            self._finish_episode(success, failure_reason)
        self.logger.info(
            "step episode_id=%s step=%s action=%s target_id=%s text=%r translated=%s reward=%.3f browser_reward=%.3f progress_score=%.3f progress_signals=%s done=%s success=%s failure_reason=%s oracle_calls=%s invalid_actions=%s repeated_actions=%s",
            self._state.episode_id,
            self._state.step_count,
            action_to_execute.action_type,
            action_to_execute.target_id,
            action_to_execute.text,
            translated,
            reward_breakdown.total,
            browser_reward,
            float(progress_evidence.get("progress_score", 0.0)),
            ",".join(progress_evidence.get("signals", [])),
            done,
            success,
            failure_reason,
            self.constraints.oracle_calls,
            self.constraints.invalid_actions,
            self.constraints.repeated_actions,
        )
        return self._make_observation(
            reward=reward_breakdown.total,
            done=done,
            success=success,
            failure_reason=failure_reason,
            reward_breakdown=reward_breakdown,
            metadata={
                "event": "step",
                "runtime": "browsergym",
                "requested_action": _dump_action(action),
                "executed_action": _dump_action(action_to_execute),
                "browsergym_action": translated,
                "target_resolution": translation_metadata,
                "browsergym_reward": browser_reward,
                "terminated": terminated,
                "truncated": truncated,
                "raw_info": dict(info),
                "judge": None
                if judge_result is None
                else {
                    "label": judge_result.label,
                    "reward": judge_result.reward,
                    "rationale": judge_result.rationale,
                },
                "failure_detail": dict(failure_detail),
                "progress_evidence": dict(progress_evidence),
                "curriculum": self.current_variant.to_config(),
                "artifacts": artifact_paths,
            },
        )

    @property
    def state(self) -> State:
        return self._state

    def close(self) -> None:
        self.logger.info(
            "close episode_id=%s step=%s task_id=%s",
            self._state.episode_id,
            self._state.step_count,
            self.current_variant.task_id,
        )
        self.runtime.close()

    def _make_observation(
        self,
        *,
        reward: float,
        done: bool,
        success: bool,
        failure_reason: str,
        reward_breakdown: Any = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> BrowserObservation:
        history_window = self.history[-5:]
        variant_config = dict(self.current_variant.to_config())
        variant_config["instruction"] = extract_instruction(self.raw_obs, self.raw_info)
        variant_config["history_texts"] = history_window
        elements, observation_stats = compact_elements_with_stats(self.raw_obs, variant_config)
        self.latest_observation_stats = observation_stats
        breakdown = reward_breakdown or self.reward_engine.compute(
            browsergym_reward=0.0,
            success=False,
            step_delta=0,
            progress_delta=0,
            oracle_delta=0,
            mistake_delta=0,
            repetition_delta=0,
        )
        metadata_payload = dict(metadata or {})
        metadata_payload["observation_stats"] = observation_stats
        metadata_payload.setdefault("observation_filter", "heuristic_ranker_v1")
        return BrowserObservation(
            episode_id=self._state.episode_id or "",
            task_id=self.current_variant.task_id,
            task_family=self.current_variant.task_family,
            difficulty=self.current_variant.difficulty,
            instruction=extract_instruction(self.raw_obs, self.raw_info),
            url=extract_url(self.raw_obs),
            step_index=self._state.step_count,
            max_steps=self.current_variant.max_steps,
            elements=elements,
            history=history_window,
            constraints=self.constraints,
            reward_breakdown=breakdown,
            done=done,
            reward=reward,
            success=success,
            failure_reason=failure_reason,  # type: ignore[arg-type]
            metadata=metadata_payload,
        )

    def _new_constraints(self) -> ConstraintState:
        llm_budget = 3
        for env_name in ("BROWSER_ENV_ORACLE_BUDGET", "BROWSER_ENV_LLM_BUDGET"):
            raw_value = os.getenv(env_name)
            if not raw_value:
                continue
            try:
                parsed = int(raw_value)
            except ValueError:
                continue
            if parsed > 0:
                llm_budget = parsed
                break
        return ConstraintState(
            step_budget=self.current_variant.max_steps,
            llm_budget=llm_budget,
            oracle_calls=0,
            invalid_actions=0,
            repeated_actions=0,
            delayed_failures=0,
            current_difficulty=self.current_variant.difficulty,
            curriculum_variant_id=self.current_variant.variant_id,
        )

    def _penalized_invalid_observation(
        self,
        action: BrowserAction,
        failure_reason: str,
        *,
        oracle_delta: int = 0,
    ) -> BrowserObservation:
        self.constraints.invalid_actions += 1
        self._state.step_count += 1
        terminal = (
            failure_reason == "oracle_budget_exceeded"
            or self._state.step_count >= self.current_variant.max_steps
            or self.constraints.invalid_actions >= self.max_invalid_actions
        )
        if failure_reason == "invalid_action" and self.constraints.invalid_actions >= self.max_invalid_actions:
            failure_reason = "too_many_invalid_actions"
        breakdown = self.reward_engine.compute(
            browsergym_reward=0.0,
            success=False,
            step_delta=1,
            progress_delta=0,
            oracle_delta=oracle_delta,
            mistake_delta=1,
            repetition_delta=0,
        )
        self.history.append(
            {
                "action_type": action.action_type,
                "target_id": action.target_id,
                "text": action.text,
                "browsergym_action": None,
                "reward": breakdown.total,
                "status": failure_reason,
            }
        )
        self.trajectory.append(
            {
                "step": self._state.step_count,
                "action": _dump_action(action),
                "executed_action": _dump_action(action),
                "browsergym_action": None,
                "browsergym_reward": 0.0,
                "reward_breakdown": _dump_model(breakdown),
                "success": False,
                "failure_reason": failure_reason,
                "failure_detail": {
                    "mode": failure_reason,
                    "action_type": action.action_type,
                    "target_id": action.target_id,
                    "translated_action": None,
                },
                "constraints": _dump_model(self.constraints),
            }
        )
        self.total_reward += breakdown.total
        if terminal:
            self._finish_episode(False, failure_reason)
        self.logger.warning(
            "invalid_observation episode_id=%s step=%s action=%s target_id=%s failure_reason=%s terminal=%s",
            self._state.episode_id,
            self._state.step_count,
            action.action_type,
            action.target_id,
            failure_reason,
            terminal,
        )
        return self._make_observation(
            reward=breakdown.total,
            done=terminal,
            success=False,
            failure_reason=failure_reason,
            reward_breakdown=breakdown,
            metadata={
                "event": "invalid_action",
                "action": _dump_action(action),
                "failure_detail": {
                    "mode": failure_reason,
                    "action_type": action.action_type,
                    "target_id": action.target_id,
                    "translated_action": None,
                },
            },
        )

    def _finish_episode(self, success: bool, failure_reason: str) -> None:
        if self.episode_done:
            return
        self.episode_done = True
        self.final_success = success
        self.final_failure_reason = failure_reason
        self.final_failure_detail = dict(self.trajectory[-1].get("failure_detail") or {}) if self.trajectory else {}
        finished_variant_id = self.current_variant.variant_id
        finished_variant_config = self.current_variant.to_config()
        episode = {
            "episode_id": self._state.episode_id,
            "task_id": self.current_variant.task_id,
            "variant_id": self.current_variant.variant_id,
            "seed": self.current_variant.seed,
            "steps": self._state.step_count,
            "success": success,
            "failure_reason": failure_reason,
            "failure_detail": dict(self.final_failure_detail),
            "progress_evidence": dict(self.last_progress_evidence),
            "total_reward": self.total_reward,
            "oracle_calls": self.constraints.oracle_calls,
            "invalid_actions": self.constraints.invalid_actions,
            "trajectory": self.trajectory,
        }
        self.replay.record_episode(episode)
        self.replay.flush()
        self.logger.info(
            "episode_finished episode_id=%s task_id=%s variant=%s success=%s failure_reason=%s steps=%s total_reward=%.3f oracle_calls=%s invalid_actions=%s",
            self._state.episode_id,
            self.current_variant.task_id,
            self.current_variant.variant_id,
            success,
            failure_reason,
            self._state.step_count,
            self.total_reward,
            self.constraints.oracle_calls,
            self.constraints.invalid_actions,
        )
        if success:
            self.curriculum.note_success(finished_variant_id)
            self.curriculum.advance_on_success()
        else:
            count = self.curriculum.note_failure(finished_variant_id)
            if self.curriculum_advance_on_failure_after and count >= self.curriculum_advance_on_failure_after:
                next_variant = self.curriculum.advance_on_failure()
                self.logger.info(
                    "curriculum_advance_on_failure episode_id=%s from_variant=%s to_variant=%s failure_count=%s threshold=%s",
                    self._state.episode_id,
                    finished_variant_id,
                    next_variant.variant_id,
                    count,
                    self.curriculum_advance_on_failure_after,
                )
            if count >= 2 and self.task_generator_enabled:
                try:
                    mutation = self.agents.mutate_task(
                        finished_variant_config,
                        {"reason": failure_reason, "count": count},
                    )
                except RuntimeError:
                    mutation = None
                if mutation and self._variant_is_valid(mutation):
                    before_count = len(self.curriculum.variants)
                    validated = self.curriculum.add_validated_variant(mutation)
                    if len(self.curriculum.variants) > before_count:
                        self.logger.info(
                            "curriculum_mutation_added episode_id=%s base_variant=%s mutated_variant=%s task_id=%s",
                            self._state.episode_id,
                            finished_variant_id,
                            validated.variant_id,
                            validated.task_id,
                        )
                    else:
                        self.logger.info(
                            "curriculum_mutation_duplicate episode_id=%s base_variant=%s variant=%s task_id=%s",
                            self._state.episode_id,
                            finished_variant_id,
                            validated.variant_id,
                            validated.task_id,
                        )

    def _target_is_valid(self, action: BrowserAction, elements: List[Any]) -> bool:
        if action.action_type in {"noop", "scroll"}:
            return True
        if action.action_type == "ask_oracle":
            return True
        ids = {element.id for element in elements if not element.id.startswith("distractor")}
        if action.action_type == "submit":
            return not action.target_id or action.target_id in ids
        if not action.target_id:
            return False
        return action.target_id in ids

    def _repair_oracle_action(
        self,
        suggested: BrowserAction,
        elements: List[Any],
        *,
        instruction: str,
    ) -> BrowserAction:
        """Apply conservative guardrails to avoid oracle getting stuck in type loops."""
        submit_after_select = self._maybe_submit_after_recent_select(suggested, elements, instruction=instruction)
        if submit_after_select is not None:
            return submit_after_select
        if suggested.action_type == "select":
            return self._repair_select_oracle_action(suggested, elements, instruction=instruction)
        if suggested.action_type != "type":
            return self._repair_non_type_oracle_action(suggested, elements, instruction=instruction)
        if not suggested.target_id:
            return suggested
        typed_text = (suggested.text or "").strip()
        if not typed_text:
            return suggested

        repeat_count = self._recent_identical_type_count(suggested.target_id)
        text_present = self._elements_contain_text(elements, typed_text)
        if repeat_count < self.oracle_submit_after_type_repeat and not text_present:
            return suggested

        submit_target = self._find_clickable_submit_target(elements)
        if submit_target:
            self.logger.info(
                "oracle_auto_submit episode_id=%s step=%s from_target=%s submit_target=%s repeat_count=%s text_present=%s",
                self._state.episode_id,
                self._state.step_count,
                suggested.target_id,
                submit_target,
                repeat_count,
                text_present,
            )
            return BrowserAction(
                action_type="click",
                target_id=submit_target,
                reasoning="auto_submit_after_repeated_type",
            )

        self.logger.info(
            "oracle_auto_submit episode_id=%s step=%s from_target=%s submit_target=%s repeat_count=%s text_present=%s",
            self._state.episode_id,
            self._state.step_count,
            suggested.target_id,
            None,
            repeat_count,
            text_present,
        )
        return BrowserAction(action_type="submit", reasoning="auto_submit_after_repeated_type")

    def _maybe_submit_after_recent_select(
        self,
        suggested: BrowserAction,
        elements: List[Any],
        *,
        instruction: str,
    ) -> Optional[BrowserAction]:
        """Turn select loops into submit actions for select-then-submit tasks."""

        if not self.history:
            return None
        last = self.history[-1]
        if last.get("action_type") != "select":
            return None
        if suggested.action_type not in {"select", "noop", "scroll"}:
            return None

        instruction_lc = instruction.lower()
        submit_target = self._find_clickable_submit_target(elements)
        needs_submit = bool(submit_target) or "submit" in instruction_lc or "click submit" in instruction_lc
        if not needs_submit:
            return None

        if submit_target:
            self.logger.info(
                "oracle_auto_submit_after_recent_select episode_id=%s step=%s submit_target=%s suggested_action=%s",
                self._state.episode_id,
                self._state.step_count,
                submit_target,
                suggested.action_type,
            )
            return BrowserAction(
                action_type="click",
                target_id=submit_target,
                reasoning="auto_submit_after_recent_select",
            )
        self.logger.info(
            "oracle_auto_submit_after_recent_select episode_id=%s step=%s submit_target=%s suggested_action=%s",
            self._state.episode_id,
            self._state.step_count,
            None,
            suggested.action_type,
        )
        return BrowserAction(action_type="submit", reasoning="auto_submit_after_recent_select")

    def _repair_select_oracle_action(
        self,
        suggested: BrowserAction,
        elements: List[Any],
        *,
        instruction: str,
    ) -> BrowserAction:
        target_phrase = self._extract_instruction_target_phrase(instruction)
        repaired = suggested
        cleaned_text = (suggested.text or "").strip()
        low_signal_text = cleaned_text.lower() in {"", "option", "select", "choose", "item", "test"}
        if target_phrase and (low_signal_text or cleaned_text.lower() != target_phrase):
            repaired = BrowserAction(
                action_type="select",
                target_id=suggested.target_id,
                text=target_phrase,
                reasoning="auto_select_text_from_instruction",
            )
            cleaned_text = target_phrase

        repeat_count = self._recent_identical_action_count(
            "select",
            repaired.target_id,
            cleaned_text,
        )
        if repeat_count < self.oracle_submit_after_select_repeat:
            return self._repair_non_type_oracle_action(repaired, elements, instruction=instruction)

        submit_target = self._find_clickable_submit_target(elements)
        if submit_target:
            self.logger.info(
                "oracle_auto_submit_after_select episode_id=%s step=%s from_target=%s submit_target=%s repeat_count=%s",
                self._state.episode_id,
                self._state.step_count,
                repaired.target_id,
                submit_target,
                repeat_count,
            )
            return BrowserAction(
                action_type="click",
                target_id=submit_target,
                reasoning="auto_submit_after_repeated_select",
            )
        self.logger.info(
            "oracle_auto_submit_after_select episode_id=%s step=%s from_target=%s submit_target=%s repeat_count=%s",
            self._state.episode_id,
            self._state.step_count,
            repaired.target_id,
            None,
            repeat_count,
        )
        return BrowserAction(action_type="submit", reasoning="auto_submit_after_repeated_select")

    def _repair_non_type_oracle_action(
        self,
        suggested: BrowserAction,
        elements: List[Any],
        *,
        instruction: str,
    ) -> BrowserAction:
        # Handle repeated click/select loops by switching to a better candidate that
        # matches instruction text or a submit control.
        if suggested.action_type not in {"click", "select", "submit"}:
            return suggested

        repeat_count = self._recent_identical_action_count(
            suggested.action_type,
            suggested.target_id,
            suggested.text,
        )
        if repeat_count <= 0 and suggested.action_type != "submit":
            return suggested

        target_phrase = self._extract_instruction_target_phrase(instruction)
        tried_ids = self._recent_target_ids(window=4)
        if suggested.target_id:
            tried_ids.add(suggested.target_id)

        candidate = self._pick_best_click_target(
            elements,
            target_phrase=target_phrase,
            exclude_ids=tried_ids,
        )
        if candidate and candidate != suggested.target_id:
            self.logger.info(
                "oracle_auto_retarget episode_id=%s step=%s action=%s from_target=%s to_target=%s repeat_count=%s phrase=%r",
                self._state.episode_id,
                self._state.step_count,
                suggested.action_type,
                suggested.target_id,
                candidate,
                repeat_count,
                target_phrase,
            )
            return BrowserAction(
                action_type="click",
                target_id=candidate,
                reasoning="auto_retarget_after_repeated_click_or_select",
            )

        if suggested.action_type == "submit":
            return suggested
        return suggested

    def _recent_identical_type_count(self, target_id: str) -> int:
        count = 0
        for item in reversed(self.history):
            if item.get("action_type") != "type":
                break
            if item.get("target_id") != target_id:
                break
            count += 1
        return count

    def _recent_identical_action_count(self, action_type: str, target_id: Optional[str], text: Optional[str]) -> int:
        count = 0
        wanted_text = (text or "").strip()
        ignore_text = action_type in {"click", "select"}
        for item in reversed(self.history):
            if item.get("action_type") != action_type:
                break
            if (item.get("target_id") or None) != (target_id or None):
                break
            if not ignore_text and (item.get("text") or "").strip() != wanted_text:
                break
            count += 1
        return count

    def _recent_target_ids(self, window: int = 4) -> set[str]:
        out: set[str] = set()
        for item in self.history[-max(0, window) :]:
            target_id = item.get("target_id")
            if target_id:
                out.add(str(target_id))
        return out

    def _elements_contain_text(self, elements: List[Any], text: str) -> bool:
        wanted = text.strip().lower()
        if not wanted:
            return False
        for element in elements:
            value = (getattr(element, "text", "") or "").strip().lower()
            if value == wanted:
                return True
        return False

    def _find_clickable_submit_target(self, elements: List[Any]) -> Optional[str]:
        keywords = ("submit", "send", "ok", "done", "go", "next", "continue")
        for element in elements:
            attrs = getattr(element, "attributes", {}) or {}
            clickable = bool(attrs.get("clickable", False))
            if not clickable:
                continue
            if not bool(getattr(element, "visible", True)):
                continue
            if not bool(getattr(element, "enabled", True)):
                continue
            haystack = " ".join(
                [
                    str(getattr(element, "text", "") or "").lower(),
                    str(getattr(element, "role", "") or "").lower(),
                    str(getattr(element, "tag", "") or "").lower(),
                    str(attrs.get("name", "")).lower(),
                    str(attrs.get("aria-label", "")).lower(),
                ]
            )
            if any(keyword in haystack for keyword in keywords):
                return str(getattr(element, "id", "") or "")
        return None

    def _extract_instruction_target_phrase(self, instruction: str) -> Optional[str]:
        if not instruction:
            return None
        quoted = re.findall(r'"([^"]+)"|\'([^\']+)\'', instruction)
        for pair in quoted:
            value = (pair[0] or pair[1] or "").strip()
            if value:
                return value.lower()
        return None

    def _pick_best_click_target(
        self,
        elements: List[Any],
        *,
        target_phrase: Optional[str],
        exclude_ids: set[str],
    ) -> Optional[str]:
        scored: List[Tuple[Tuple[int, int, int, int], str]] = []
        for element in elements:
            element_id = str(getattr(element, "id", "") or "")
            if not element_id or element_id in exclude_ids:
                continue
            attrs = getattr(element, "attributes", {}) or {}
            clickable = bool(attrs.get("clickable", False))
            if not clickable:
                continue
            visible = bool(getattr(element, "visible", True))
            enabled = bool(getattr(element, "enabled", True))
            text = str(getattr(element, "text", "") or "").strip().lower()
            role = str(getattr(element, "role", "") or "").strip().lower()
            phrase_score = 0
            if target_phrase and text:
                if text == target_phrase:
                    phrase_score = 3
                elif target_phrase in text:
                    phrase_score = 2
                elif text in target_phrase:
                    phrase_score = 1
            submit_score = int("submit" in text or "submit" in role)
            score = (phrase_score, submit_score, int(visible), int(enabled))
            scored.append((score, element_id))
        if not scored:
            return None
        scored.sort(reverse=True)
        return scored[0][1]

    def _has_form_action(self) -> bool:
        return any(item.get("action_type") in {"type", "select"} for item in self.history)

    def _is_submit_like_action(self, action: BrowserAction, elements: List[Any]) -> bool:
        """Return whether an action behaves like a form submission attempt."""
        if action.action_type == "submit":
            return True
        if action.action_type != "click" or not action.target_id:
            return False
        submit_target = self._find_clickable_submit_target(elements)
        return bool(submit_target and action.target_id == submit_target)

    def _is_done(self) -> bool:
        return self.episode_done or self._state.step_count >= self.current_variant.max_steps

    def _build_failure_detail(
        self,
        *,
        action: BrowserAction,
        translated: str,
        browser_reward: float,
        terminated: bool,
        truncated: bool,
        submit_like_failure: bool,
        mode: str,
    ) -> Dict[str, Any]:
        detail_mode = mode
        if truncated:
            detail_mode = "runtime_truncated"
        elif mode == "submission_failed" and submit_like_failure and not self._has_form_action():
            detail_mode = "premature_submit"
        elif mode == "submission_failed" and submit_like_failure and self._has_form_action():
            detail_mode = "submit_without_success"
        elif mode == "task_failed" and terminated and browser_reward <= 0.0:
            detail_mode = "task_failed_terminal"
        elif terminated and browser_reward <= 0.0:
            detail_mode = "terminal_without_success"
        return {
            "mode": detail_mode,
            "base_reason": mode,
            "action_type": action.action_type,
            "target_id": action.target_id,
            "translated_action": translated,
            "browser_reward": browser_reward,
            "terminated": terminated,
            "truncated": truncated,
            "has_form_action": self._has_form_action(),
            "no_progress_streak": self.consecutive_no_progress_steps,
            "progress_score": float(self.last_progress_evidence.get("progress_score", 0.0)),
            "progress_signals": list(self.last_progress_evidence.get("signals", [])),
        }

    def _delayed_penalty_for_failure(
        self,
        *,
        failure_reason: str,
        failure_detail: Dict[str, Any],
    ) -> float:
        if failure_reason not in {"submission_failed", "low_progress_abort", "task_failed"}:
            return 0.0
        mode = str(failure_detail.get("mode") or "")
        if mode == "premature_submit":
            return -1.25
        if mode == "submit_without_success":
            return -1.0
        if mode == "runtime_truncated":
            return -0.5
        if mode == "terminal_without_success":
            return -0.75
        if mode == "task_failed_terminal":
            return -0.75
        if self._has_form_action():
            return -1.0
        return -0.5

    def _observation_config(self) -> Dict[str, Any]:
        variant_config = dict(self.current_variant.to_config())
        variant_config["instruction"] = extract_instruction(self.raw_obs, self.raw_info)
        variant_config["history_texts"] = self.history[-5:]
        return variant_config

    def _current_elements_with_stats(self) -> Tuple[List[Any], Dict[str, Any]]:
        return compact_elements_with_stats(self.raw_obs, self._observation_config())

    def _element_state_signature(self, element: Any) -> Tuple[Any, ...]:
        attrs = dict(getattr(element, "attributes", {}) or {})
        return (
            getattr(element, "text", "") or "",
            attrs.get("value"),
            attrs.get("aria-selected"),
            attrs.get("aria-expanded"),
            attrs.get("aria-pressed"),
            attrs.get("checked"),
            attrs.get("open"),
            attrs.get("selected"),
            attrs.get("class"),
            getattr(element, "visible", True),
            getattr(element, "enabled", True),
        )

    def _visible_submit_targets(self, elements: List[Any]) -> List[str]:
        targets: List[str] = []
        for element in elements:
            attrs = dict(getattr(element, "attributes", {}) or {})
            hints = attrs.get("semantic_hints") or []
            if not getattr(element, "visible", True) or not getattr(element, "enabled", True):
                continue
            if "submit_like" in hints:
                targets.append(str(getattr(element, "id", "")))
        return targets

    def _assess_progress_evidence(
        self,
        *,
        previous_elements: List[Any],
        action: BrowserAction,
        browser_reward: float,
        success: bool,
    ) -> Dict[str, Any]:
        current_elements, current_stats = self._current_elements_with_stats()
        previous_map = {str(getattr(element, "id", "")): element for element in previous_elements}
        current_map = {str(getattr(element, "id", "")): element for element in current_elements}
        previous_visible_ids = {element_id for element_id, element in previous_map.items() if getattr(element, "visible", True)}
        current_visible_ids = {element_id for element_id, element in current_map.items() if getattr(element, "visible", True)}
        union_ids = previous_visible_ids | current_visible_ids
        changed_ratio = 0.0
        if union_ids:
            changed_ratio = len(previous_visible_ids ^ current_visible_ids) / float(len(union_ids))

        target_id = str(action.target_id or "")
        recent_targets = [
            str(item.get("target_id") or "")
            for item in self.history[-3:]
            if item.get("target_id")
        ]
        target_changed_from_recent_history = bool(target_id and target_id not in recent_targets)

        previous_target = previous_map.get(target_id) if target_id else None
        current_target = current_map.get(target_id) if target_id else None
        previous_signature = self._element_state_signature(previous_target) if previous_target is not None else None
        current_signature = self._element_state_signature(current_target) if current_target is not None else None
        target_state_changed = (
            previous_signature is not None
            and current_signature is not None
            and previous_signature != current_signature
        )
        text_entered = action.action_type in {"type", "clear"} and target_state_changed
        selection_changed = action.action_type == "select" and target_state_changed
        click_target_state_changed = action.action_type == "click" and target_state_changed

        previous_submit_targets = set(self._visible_submit_targets(previous_elements))
        current_submit_targets = set(self._visible_submit_targets(current_elements))
        submit_target_appeared = bool(current_submit_targets - previous_submit_targets)

        previous_top_ids = [str(getattr(element, "id", "")) for element in previous_elements[:3]]
        current_top_ids = [str(getattr(element, "id", "")) for element in current_elements[:3]]
        new_high_rank_element_surfaced = bool([element_id for element_id in current_top_ids if element_id and element_id not in previous_top_ids])

        material_element_set_change = changed_ratio >= 0.25 or len(previous_visible_ids ^ current_visible_ids) >= 2
        positive_browser_reward = browser_reward > 0.0

        signal_weights = {
            "positive_browser_reward": 2.0,
            "material_element_set_change": 1.5,
            "text_entered": 1.25,
            "selection_changed": 1.25,
            "click_target_state_changed": 1.15,
            "submit_target_appeared": 1.0,
            "new_high_rank_element_surfaced": 0.9,
            "target_changed_from_recent_history": 0.35,
        }
        raw_signals = {
            "positive_browser_reward": positive_browser_reward,
            "material_element_set_change": material_element_set_change,
            "text_entered": text_entered,
            "selection_changed": selection_changed,
            "click_target_state_changed": click_target_state_changed,
            "submit_target_appeared": submit_target_appeared,
            "new_high_rank_element_surfaced": new_high_rank_element_surfaced,
            "target_changed_from_recent_history": target_changed_from_recent_history,
        }
        progress_score = sum(signal_weights[name] for name, active in raw_signals.items() if active)
        strong_signal = any(
            raw_signals[name]
            for name in (
                "positive_browser_reward",
                "material_element_set_change",
                "text_entered",
                "selection_changed",
                "click_target_state_changed",
                "submit_target_appeared",
                "new_high_rank_element_surfaced",
            )
        )
        progress_signal = bool(success or progress_score >= 0.9)
        reward_credit = int(
            strong_signal
            and not success
            and browser_reward <= 0.0
        )
        return {
            "progress_signal": progress_signal,
            "reward_credit": reward_credit,
            "progress_score": round(progress_score, 4),
            "signals": [name for name, active in raw_signals.items() if active],
            "changed_visible_ratio": round(changed_ratio, 4),
            "previous_visible_count": len(previous_visible_ids),
            "current_visible_count": len(current_visible_ids),
            "previous_top_ids": previous_top_ids,
            "current_top_ids": current_top_ids,
            "recent_target_ids": recent_targets,
            "current_observation_stats": current_stats,
        }

    def _variant_is_valid(self, config: Dict[str, Any]) -> bool:
        try:
            variant = TaskVariant.from_config(config)
        except Exception:
            return False

        if not variant.task_id:
            return False
        if int(variant.max_steps) <= 0:
            return False

        try:
            import gymnasium as gym
            import browsergym.miniwob  # noqa: F401
        except Exception:
            return False

        if variant.task_id not in {str(key) for key in gym.registry.keys()}:
            return False
        if "miniwob" in variant.task_id and not os.getenv("MINIWOB_URL"):
            return False

        try:
            headless = os.getenv("BROWSER_ENV_HEADLESS", "1") != "0"
            env = gym.make(variant.task_id, headless=headless)
            try:
                env.reset(seed=variant.seed)
            finally:
                env.close()
        except Exception:
            return False
        return True


def _ensure_mapping(value: Any) -> Mapping[str, Any]:
    return value if isinstance(value, Mapping) else {}


def _dump_model(model: Any) -> Dict[str, Any]:
    if hasattr(model, "model_dump"):
        return model.model_dump()
    if hasattr(model, "dict"):
        return model.dict()
    return dict(model)


def _dump_action(action: BrowserAction) -> Dict[str, Any]:
    return _dump_model(action)