"""TalkingHeadBench OpenEnv Environment. Wraps the existing benchmark into a step-based environment with three agent-facing episode types: 1) image audit: reset(image) -> node1 -> node2 -> done (reward=subenv1) 2) clip audit: reset(clips) -> node5 -> done (reward=subenv2) 3) weight audit: reset(weights) -> node8 -> done (reward=subenv3) 4) clips+weights: reset(clips_and_weights) -> node5 -> node8 -> done The environment supports two modes: - benchmark mode: pulls deterministic cases from tests/test_set - custom mode: consumes user-ingested signal bundles (via ingestion_id) """ from __future__ import annotations import json import logging import random from pathlib import Path from typing import Any, Optional, TypeVar from uuid import uuid4 from models import TalkingHeadObservation from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import Action, State from pydantic import ValidationError from server.artifact_ingest import get_ingested_bundle from src.envs.subenv1.node1_image_diagnostician import diagnose_image from src.envs.subenv1.node2_param_anomaly import detect_param_anomalies from src.envs.subenv1.node3_grader import ( grade_anomaly_detection, grade_image_diagnostics, produce_reference_audit_handoff, ) from src.envs.subenv2.node5_disposition import recommend_clip_disposition from src.envs.subenv2.node6_grader import grade_clip_disposition from src.envs.subenv3.node8_phoneme_risk import assess_phoneme_risk from src.envs.subenv3.node9_grader import grade_behavioral_audit from src.schemas.ground_truth import ( GroundTruthBehavioralAnnotation, GroundTruthClipAnnotation, GroundTruthImageAnnotation, GroundTruthParamAnnotation, ) from src.schemas.subenv1 import ( ImageDiagnosticsAction, ImageDiagnosticsObservation, ParamAnomalyAction, ParamAnomalyObservation, ) from src.schemas.subenv2 import ( ClipDispositionAction, ClipDispositionObservation, ClipEvidenceDossier, ClipSignalObservation, ) from src.schemas.subenv3 import ( LayerAnomalyFlag, PhonemeRiskAction, PhonemeRiskObservation, TokenAnomalyFlag, WeightEvidenceDossier, WeightSignalObservation, ) log = logging.getLogger(__name__) _T = TypeVar("_T") _TEST_SET_DIR = Path(__file__).resolve().parent.parent / "tests" / "test_set" API_VERSION = "1.0" class TalkingHeadEnvironment(Environment[Action, TalkingHeadObservation, State]): """OpenEnv adapter for TalkingHeadBench.""" def __init__(self) -> None: self._state = State(episode_id=str(uuid4()), step_count=0) self._subenv1_cases = self._load_cases("subenv1_cases.json") self._subenv2_cases = self._load_cases("subenv2_cases.json") self._subenv3_cases = self._load_cases("subenv3_cases.json") self._cursor = 0 self._phase = "idle" self._scores: dict[str, float] = {} self._current_case: dict[str, Any] | None = None self._node1_obs: ImageDiagnosticsObservation | None = None self._node1_gt: GroundTruthImageAnnotation | None = None self._node1_action: ImageDiagnosticsAction | None = None self._node2_obs: ParamAnomalyObservation | None = None self._node2_gt: GroundTruthParamAnnotation | None = None self._clip_obs: ClipSignalObservation | None = None self._clip_gt: GroundTruthClipAnnotation | None = None self._subenv3_case: dict[str, Any] | None = None self._subenv3_gt: GroundTruthBehavioralAnnotation | None = None self._custom_clip_obs: list[ClipSignalObservation] | None = None self._custom_weight_obs: WeightSignalObservation | None = None self._proposed_config: dict[str, float] = self._default_param_config() self._mode: str = "benchmark" self._episode_mode: str = "image" self._provenance: dict[str, Any] = {} def _load_cases(self, filename: str) -> list[dict[str, Any]]: path = _TEST_SET_DIR / filename if not path.exists(): log.warning("Missing test case file: %s", path) return [] with path.open("r", encoding="utf-8") as f: payload = json.load(f) cases = payload.get("cases", []) return cases if isinstance(cases, list) else [] def _default_param_config(self) -> dict[str, float]: return {"cfg": 7.5, "denoise_alt": 0.30, "eta": 0.10} def _reset_runtime_state(self, episode_id: Optional[str]) -> None: self._state = State(episode_id=episode_id or str(uuid4()), step_count=0) self._scores = {} self._phase = "idle" self._mode = "benchmark" self._episode_mode = "image" self._provenance = { "bundle_source": "benchmark_test_set", "ground_truth_source": "annotations", } self._node1_action = None self._node2_obs = None self._node2_gt = None self._clip_obs = None self._clip_gt = None self._subenv3_case = None self._subenv3_gt = None self._custom_clip_obs = None self._custom_weight_obs = None self._proposed_config = self._default_param_config() def _current_case_id(self) -> str: if self._current_case and "id" in self._current_case: return str(self._current_case.get("id")) return "unknown" def _build_agent_observation( self, *, done: bool, reward: float, node: str, step: int, instruction: Optional[str], expected_action_schema: Optional[str], signals: dict[str, Any], scores: Optional[dict[str, Any]] = None, reward_formula: Optional[str] = None, error: Optional[str] = None, ) -> TalkingHeadObservation: return TalkingHeadObservation( done=done, reward=reward, node=node, step=step, api_version=API_VERSION, mode="custom" if self._mode == "custom" else "benchmark", is_deterministic=True, case_id=self._current_case_id(), episode_id=self._state.episode_id, instruction=instruction, expected_action_schema=expected_action_schema, signals=signals, provenance=dict(self._provenance), scores=scores, reward_formula=reward_formula, error=error, ) def _normalize_episode_mode(self, mode: Optional[str]) -> str: value = (mode or "image").strip().lower() valid = {"image", "clips", "weights", "clips_and_weights"} if value not in valid: raise RuntimeError( "mode must be one of: image, clips, weights, clips_and_weights" ) return value def _next_index(self, total: int, seed: Optional[int]) -> int: if total <= 0: raise RuntimeError("Cannot sample from an empty case list") if seed is not None: return random.Random(seed).randrange(total) idx = self._cursor % total self._cursor += 1 return idx def _infer_mode_from_custom_bundle(self) -> str: has_image = self._node1_obs is not None has_clips = bool(self._custom_clip_obs) has_weight = self._custom_weight_obs is not None or self._subenv3_case is not None if has_clips and has_weight: return "clips_and_weights" if has_clips: return "clips" if has_weight: return "weights" if has_image: return "image" raise RuntimeError("Custom bundle does not include image, clips, or weight signals.") def _build_clip_disposition_observation( self, clip_obs: ClipSignalObservation, ) -> ClipDispositionObservation: dossier = self._build_clip_evidence_dossier(clip_obs) damage_map = {"none": 0.0, "minor": 0.05, "moderate": 0.20, "severe": 0.50} marginal_damage = damage_map.get(dossier.identity_drift_severity, 0.1) return ClipDispositionObservation( evidence_dossier=dossier, minimum_clips_needed=20, phoneme_gap_severity={}, pose_gap_severity={}, budget_remaining=max(0, 50 - clip_obs.clips_audited_so_far), marginal_training_damage=marginal_damage, marginal_coverage_gain=float(clip_obs.phoneme_coverage_new), ) def _start_image_episode(self, seed: Optional[int]) -> TalkingHeadObservation: if self._mode == "custom": if self._node1_obs is None: raise RuntimeError("Custom bundle does not include image_observation.") self._phase = "node1" return self._build_agent_observation( done=False, reward=0.0, node="node1_image_diagnostician", step=0, instruction=( "Diagnose reference-image quality and prompt risks from the provided " "signals." ), expected_action_schema="ImageDiagnosticsAction", signals=self._node1_obs.model_dump(mode="json"), ) if not self._subenv1_cases: raise RuntimeError("subenv1_cases.json has no cases; cannot reset environment") idx = self._next_index(len(self._subenv1_cases), seed) self._current_case = self._subenv1_cases[idx] self._node1_obs = ImageDiagnosticsObservation.model_validate( self._current_case["observation"] ) self._node1_gt = GroundTruthImageAnnotation.model_validate( self._current_case["ground_truth"] ) self._phase = "node1" return self._build_agent_observation( done=False, reward=0.0, node="node1_image_diagnostician", step=0, instruction=( "Diagnose reference-image quality and prompt risks from the provided " "signals." ), expected_action_schema="ImageDiagnosticsAction", signals=self._node1_obs.model_dump(mode="json"), ) def _start_clips_episode(self, seed: Optional[int]) -> TalkingHeadObservation: if self._mode == "custom": if not self._custom_clip_obs: raise RuntimeError("Custom bundle does not include clip_signal_observations.") self._clip_obs = self._custom_clip_obs[0] clip_obs_payload = self._build_clip_disposition_observation(self._clip_obs) baseline_action = recommend_clip_disposition(clip_obs_payload) self._clip_gt = self._synthesize_clip_ground_truth(baseline_action) else: if not self._subenv2_cases: raise RuntimeError("subenv2_cases.json has no cases; cannot reset environment") idx = self._next_index(len(self._subenv2_cases), seed) clip_case = self._subenv2_cases[idx] self._current_case = {"id": clip_case.get("id", f"subenv2-{idx}")} self._clip_obs = ClipSignalObservation.model_validate(clip_case["observation"]) self._clip_gt = GroundTruthClipAnnotation.model_validate(clip_case["ground_truth"]) clip_obs_payload = self._build_clip_disposition_observation(self._clip_obs) self._phase = "clips" return self._build_agent_observation( done=False, reward=0.0, node="node5_clip_disposition_recommender", step=0, instruction=( "Recommend clip disposition (accept/reject/fix/defer) from the dossier " "and dataset context." ), expected_action_schema="ClipDispositionAction", signals=clip_obs_payload.model_dump(mode="json"), ) def _start_weights_episode(self, seed: Optional[int]) -> TalkingHeadObservation: if self._mode == "custom": if self._custom_weight_obs is None and self._subenv3_case is None: raise RuntimeError("Custom bundle does not include weight_observation.") phoneme_obs = self._build_phoneme_risk_observation() if self._subenv3_gt is None: baseline = assess_phoneme_risk(phoneme_obs) self._subenv3_gt = self._synthesize_behavioral_ground_truth(baseline) else: if not self._subenv3_cases: raise RuntimeError("subenv3_cases.json has no cases; cannot reset environment") idx = self._next_index(len(self._subenv3_cases), seed) self._subenv3_case = self._pick_subenv3_case(idx) self._subenv3_gt = self._resolve_behavioral_ground_truth(self._subenv3_case) self._current_case = self._subenv3_case phoneme_obs = self._build_phoneme_risk_observation() self._phase = "node8" return self._build_agent_observation( done=False, reward=0.0, node="node8_phoneme_risk_assessor", step=0, instruction=( "Assess phoneme-level behavioral risk from weight evidence and provide " "mitigation guidance." ), expected_action_schema="PhonemeRiskAction", signals=phoneme_obs.model_dump(mode="json"), ) def _start_clips_and_weights_episode(self, seed: Optional[int]) -> TalkingHeadObservation: if self._mode == "custom": if not self._custom_clip_obs: raise RuntimeError("Custom bundle does not include clip_signal_observations.") if self._custom_weight_obs is None and self._subenv3_case is None: raise RuntimeError("Custom bundle does not include weight_observation.") self._clip_obs = self._custom_clip_obs[0] clip_obs_payload = self._build_clip_disposition_observation(self._clip_obs) baseline_action = recommend_clip_disposition(clip_obs_payload) self._clip_gt = self._synthesize_clip_ground_truth(baseline_action) else: if not self._subenv2_cases: raise RuntimeError("subenv2_cases.json has no cases; cannot reset environment") if not self._subenv3_cases: raise RuntimeError("subenv3_cases.json has no cases; cannot reset environment") idx = self._next_index(len(self._subenv2_cases), seed) clip_case = self._subenv2_cases[idx] self._subenv3_case = self._pick_subenv3_case(idx) self._subenv3_gt = self._resolve_behavioral_ground_truth(self._subenv3_case) self._current_case = {"id": clip_case.get("id", f"combo-{idx}")} self._clip_obs = ClipSignalObservation.model_validate(clip_case["observation"]) self._clip_gt = GroundTruthClipAnnotation.model_validate(clip_case["ground_truth"]) clip_obs_payload = self._build_clip_disposition_observation(self._clip_obs) self._phase = "clips" return self._build_agent_observation( done=False, reward=0.0, node="node5_clip_disposition_recommender", step=0, instruction=( "Recommend clip disposition (accept/reject/fix/defer) from the dossier " "and dataset context." ), expected_action_schema="ClipDispositionAction", signals=clip_obs_payload.model_dump(mode="json"), ) def reset( self, seed: Optional[int] = None, episode_id: Optional[str] = None, mode: Optional[str] = None, **kwargs: Any, ) -> TalkingHeadObservation: self._reset_runtime_state(episode_id) requested_mode = mode if mode is not None else kwargs.get("mode") custom_bundle = kwargs.get("custom_bundle") ingestion_id = kwargs.get("ingestion_id") if custom_bundle is None and ingestion_id is not None: custom_bundle = get_ingested_bundle(str(ingestion_id)) if custom_bundle is None: raise RuntimeError(f"Unknown ingestion_id: {ingestion_id}") if isinstance(custom_bundle, str): custom_bundle = json.loads(custom_bundle) if custom_bundle is not None: if not isinstance(custom_bundle, dict): raise RuntimeError("custom_bundle must be an object") self._mode = "custom" self._provenance = { "bundle_source": "ingestion_id" if ingestion_id is not None else "custom_bundle", "ground_truth_source": "synthesized_from_heuristics", } if ingestion_id is not None: self._provenance["ingestion_id"] = str(ingestion_id) self._initialize_custom_bundle(custom_bundle, ingestion_id) inferred_mode = self._infer_mode_from_custom_bundle() self._episode_mode = ( self._normalize_episode_mode(requested_mode) if requested_mode is not None else inferred_mode ) else: self._mode = "benchmark" self._provenance = { "bundle_source": "benchmark_test_set", "ground_truth_source": "annotations", "test_set_dir": str(_TEST_SET_DIR), } self._episode_mode = self._normalize_episode_mode(requested_mode) if self._episode_mode == "image": return self._start_image_episode(seed) if self._episode_mode == "clips": return self._start_clips_episode(seed) if self._episode_mode == "weights": return self._start_weights_episode(seed) return self._start_clips_and_weights_episode(seed) def step( self, action: Action, timeout_s: Optional[float] = None, **kwargs: Any, ) -> TalkingHeadObservation: self._state.step_count += 1 if self._phase == "node1": return self._handle_node1_action(action) if self._phase == "node2": return self._handle_node2_action(action) if self._phase == "clips": return self._handle_clips_action(action) if self._phase == "node8": return self._handle_node8_action(action) return self._build_agent_observation( done=True, reward=0.0, node="episode_complete", step=3, instruction=None, expected_action_schema=None, signals={}, error="Episode complete. Call reset() for a new episode.", ) def _handle_node1_action(self, action: Action) -> TalkingHeadObservation: if self._current_case is None or self._node1_obs is None or self._node1_gt is None: raise RuntimeError("Environment is not initialized. Call reset() before step().") node1_action = self._parse_action(action, ImageDiagnosticsAction) self._node1_action = node1_action node1_score = grade_image_diagnostics(node1_action, self._node1_gt) self._scores["node1"] = node1_score self._node2_obs = self._build_param_anomaly_observation(node1_action, self._node1_obs) self._node2_gt = self._resolve_param_ground_truth(self._current_case, self._node2_obs) self._phase = "node2" return self._build_agent_observation( done=False, reward=node1_score, node="node2_param_anomaly_detector", step=1, instruction=( "Analyze proposed generation parameters, detect anomalies, and propose " "directional fixes." ), expected_action_schema="ParamAnomalyAction", signals=self._node2_obs.model_dump(mode="json"), ) def _handle_node2_action(self, action: Action) -> TalkingHeadObservation: if ( self._node1_action is None or self._node1_gt is None or self._node2_obs is None or self._node2_gt is None ): raise RuntimeError("Node 2 prerequisites missing. Run Node 1 step first.") node2_action = self._parse_action(action, ParamAnomalyAction) node2_score = grade_anomaly_detection(node2_action, self._node2_gt) self._scores["node2"] = node2_score subenv1_score = 0.50 * self._scores.get("node1", 0.0) + 0.50 * node2_score self._scores["subenv1"] = subenv1_score # Maintain compatibility metadata for Sub-env 1 reporting only. _ = produce_reference_audit_handoff( self._node1_action, node2_action, self._node1_gt, self._node2_gt, ) self._phase = "done" return self._build_agent_observation( done=True, reward=subenv1_score, node="episode_complete", step=2, instruction=None, expected_action_schema=None, signals={}, reward_formula="subenv1_score", scores={ "subenv1_score": subenv1_score, "breakdown": self._scores, }, ) def _handle_clips_action(self, action: Action) -> TalkingHeadObservation: if self._clip_obs is None or self._clip_gt is None: raise RuntimeError("Clip phase prerequisites missing. Call reset() first.") clip_action = self._parse_action(action, ClipDispositionAction) clip_score = grade_clip_disposition(clip_action, self._clip_gt) self._scores["node5"] = clip_score self._scores["subenv2"] = clip_score if self._episode_mode == "clips": self._phase = "done" return self._build_agent_observation( done=True, reward=clip_score, node="episode_complete", step=1, instruction=None, expected_action_schema=None, signals={}, reward_formula="subenv2_score", scores={ "subenv2_score": clip_score, "breakdown": self._scores, }, ) phoneme_obs = self._build_phoneme_risk_observation() if self._subenv3_gt is None: baseline = assess_phoneme_risk(phoneme_obs) self._subenv3_gt = self._synthesize_behavioral_ground_truth(baseline) self._phase = "node8" return self._build_agent_observation( done=False, reward=clip_score, node="node8_phoneme_risk_assessor", step=1, instruction=( "Assess phoneme-level behavioral risk from weight evidence and provide " "mitigation guidance." ), expected_action_schema="PhonemeRiskAction", signals=phoneme_obs.model_dump(mode="json"), ) def _handle_node8_action(self, action: Action) -> TalkingHeadObservation: if self._subenv3_gt is None: raise RuntimeError("Behavioral ground truth is unavailable. Reset the environment.") node8_action = self._parse_action(action, PhonemeRiskAction) node8_score = grade_behavioral_audit(node8_action, self._subenv3_gt) self._scores["node8"] = node8_score self._scores["subenv3"] = node8_score s2 = self._scores.get("subenv2", 0.0) s3 = self._scores.get("subenv3", 0.0) if self._episode_mode == "clips_and_weights": final_score = 0.50 * s2 + 0.50 * s3 reward_formula = "0.50 * subenv2 + 0.50 * subenv3" scores = { "subenv2_score": s2, "subenv3_score": s3, "final_score": final_score, "breakdown": self._scores, } else: final_score = s3 reward_formula = "subenv3_score" scores = { "subenv3_score": s3, "breakdown": self._scores, } self._scores["final"] = final_score self._phase = "done" return self._build_agent_observation( done=True, reward=final_score, node="episode_complete", step=3, instruction=None, expected_action_schema=None, signals={}, reward_formula=reward_formula, scores=scores, ) def _parse_action(self, action: Action, expected_type: type[_T]) -> _T: if isinstance(action, expected_type): return action payload: Any = None if isinstance(action, dict): payload = action.get("action", action) elif isinstance(action, Action): payload = action.metadata.get("action", action.metadata) elif hasattr(action, "metadata") and isinstance(action.metadata, dict): payload = action.metadata.get("action", action.metadata) elif hasattr(action, "model_dump"): payload = action.model_dump() if isinstance(payload, str): payload = json.loads(payload) if not isinstance(payload, dict): raise ValueError( f"Unable to parse action payload for {expected_type.__name__}; " f"received {type(action).__name__}" ) return expected_type.model_validate(payload) def _initialize_custom_bundle( self, custom_bundle: dict[str, Any], ingestion_id: Optional[str], ) -> None: image_payload = ( custom_bundle.get("image_observation") or custom_bundle.get("image_obs") or custom_bundle.get("reference_image_obs") ) if isinstance(image_payload, dict): self._node1_obs = ImageDiagnosticsObservation.model_validate(image_payload) self._node1_gt = self._synthesize_image_ground_truth(self._node1_obs) else: self._node1_obs = None self._node1_gt = None config_payload = custom_bundle.get("param_config") or custom_bundle.get("proposed_config") if isinstance(config_payload, dict): for key, value in config_payload.items(): if isinstance(value, (int, float)): self._proposed_config[key] = float(value) raw_clips = ( custom_bundle.get("clip_signal_observations") or custom_bundle.get("clip_signal_obs_list") or [] ) if isinstance(raw_clips, list): self._custom_clip_obs = [ ClipSignalObservation.model_validate(item) for item in raw_clips if isinstance(item, dict) ] else: self._custom_clip_obs = [] custom_phoneme_case: dict[str, Any] | None = None raw_weight = custom_bundle.get("weight_observation") or custom_bundle.get("weight_obs") if isinstance(raw_weight, dict): try: self._custom_weight_obs = WeightSignalObservation.model_validate(raw_weight) except ValidationError: try: phoneme_obs = PhonemeRiskObservation.model_validate(raw_weight) custom_phoneme_case = {"observation": phoneme_obs.model_dump(mode="json")} except ValidationError as exc: raise RuntimeError( "weight_observation must be a WeightSignalObservation or " "PhonemeRiskObservation payload" ) from exc case_id = ( custom_bundle.get("case_id") or custom_bundle.get("id") or ingestion_id or f"custom-{uuid4().hex[:8]}" ) self._current_case = {"id": str(case_id)} metadata = custom_bundle.get("ingestion_metadata") if isinstance(metadata, dict): self._provenance["ingestion_metadata"] = metadata source_files = custom_bundle.get("source_files") if isinstance(source_files, dict): self._provenance["source_files"] = source_files self._subenv3_case = custom_phoneme_case self._subenv3_gt = None def _synthesize_image_ground_truth( self, image_obs: ImageDiagnosticsObservation, ) -> GroundTruthImageAnnotation: baseline = diagnose_image(image_obs) return GroundTruthImageAnnotation( regime_classification=baseline.regime_classification, acceptable_regimes=[baseline.regime_classification], identified_risk_factors=baseline.identified_risk_factors, valid_prompt_modifications=baseline.recommended_prompt_modifications, ) def _pick_subenv3_case(self, idx: int) -> dict[str, Any] | None: if not self._subenv3_cases: return None return self._subenv3_cases[idx % len(self._subenv3_cases)] def _resolve_behavioral_ground_truth( self, case: dict[str, Any] | None, ) -> GroundTruthBehavioralAnnotation: if case and isinstance(case.get("ground_truth"), dict): try: return GroundTruthBehavioralAnnotation.model_validate(case["ground_truth"]) except Exception as exc: # noqa: BLE001 log.warning("Failed to parse behavioral ground truth from case: %s", exc) return GroundTruthBehavioralAnnotation( phoneme_risk_ranking=[], predicted_behavior_triggers=[], risky_phoneme_clusters=[], model_behavioral_safety="safe", valid_mitigation_set=set(), ) def _build_param_anomaly_observation( self, node1_action: ImageDiagnosticsAction, image_obs: ImageDiagnosticsObservation, ) -> ParamAnomalyObservation: return ParamAnomalyObservation( proposed_config=self._proposed_config, regime=node1_action.regime_classification, identified_risk_factors=node1_action.identified_risk_factors, image_usability_score=node1_action.image_usability_score, face_occupancy_ratio=image_obs.face_occupancy_ratio, estimated_yaw_degrees=image_obs.estimated_yaw_degrees, background_complexity_score=image_obs.background_complexity_score, lighting_uniformity_score=image_obs.lighting_uniformity_score, occlusion_detected=image_obs.occlusion_detected, prompt_identity_anchoring=image_obs.identity_anchoring_strength, prompt_token_count=image_obs.prompt_token_count, conflicting_descriptors=image_obs.conflicting_descriptors, ) def _resolve_param_ground_truth( self, case: dict[str, Any], node2_obs: ParamAnomalyObservation, ) -> GroundTruthParamAnnotation: for key in ("param_ground_truth", "ground_truth_param", "node2_ground_truth"): if isinstance(case.get(key), dict): return GroundTruthParamAnnotation.model_validate(case[key]) baseline = detect_param_anomalies(node2_obs) return GroundTruthParamAnnotation( config_risk_level=baseline.config_risk_level, anomalies=baseline.anomalies, predicted_failure_modes=list(baseline.predicted_failure_modes), valid_fix_directions=baseline.directional_fixes, ) def _build_clip_evidence_dossier(self, obs: ClipSignalObservation) -> ClipEvidenceDossier: drift = obs.identity_cosine_drift if drift < 0.05: drift_severity = "none" elif drift < 0.12: drift_severity = "minor" elif drift < 0.22: drift_severity = "moderate" else: drift_severity = "severe" temporal_instability = ( obs.landmark_stability_score > 0.04 or obs.frame_difference_mean > 15.0 ) lip_sync = obs.lip_sync_confidence if lip_sync >= 0.75: lip_sync_quality = "good" elif lip_sync >= 0.50: lip_sync_quality = "acceptable" elif lip_sync >= 0.20: lip_sync_quality = "poor" else: lip_sync_quality = "absent" dataset_redundancy = min(1.0, obs.similar_clips_accepted / 10.0) if drift_severity in ("none", "minor") and not temporal_instability and lip_sync >= 0.50: training_impact = "positive" elif drift_severity == "severe" or (temporal_instability and lip_sync < 0.20): training_impact = "negative" else: training_impact = "neutral" primary_rejection_reason = None if drift_severity in ("moderate", "severe"): primary_rejection_reason = ( f"identity cosine drift {drift:.3f} exceeds threshold; severity={drift_severity}" ) elif temporal_instability and lip_sync < 0.20: primary_rejection_reason = "temporal instability with absent lip sync" summary = ( f"Drift={drift_severity}, temporal_instability={temporal_instability}, " f"lip_sync={lip_sync_quality}, phoneme_value={obs.phoneme_coverage_new:.2f}, " f"redundancy={dataset_redundancy:.2f}." ) return ClipEvidenceDossier( clip_id=obs.clip_id, identity_drift_severity=drift_severity, temporal_instability_flag=temporal_instability, lip_sync_quality=lip_sync_quality, unique_phoneme_value=obs.phoneme_coverage_new, dataset_redundancy_score=dataset_redundancy, estimated_training_impact=training_impact, primary_rejection_reason=primary_rejection_reason, evidence_summary=summary, ) def _synthesize_clip_ground_truth( self, action, ) -> GroundTruthClipAnnotation: valid_overrides = [action.override_justification] if action.override_justification else [] return GroundTruthClipAnnotation( disposition=action.disposition, confidence=float(action.confidence), disposition_ambiguity=0.6 if action.disposition == "defer" else 0.0, valid_fix_steps=list(action.fix_instructions or []), valid_override_justifications=valid_overrides, expected_reasoning_elements=[], ) def _run_subenv2_from_observations( self, clip_observations: list[ClipSignalObservation], ) -> float: if not clip_observations: return 0.5 scores: list[float] = [] for clip_obs in clip_observations: try: dossier = self._build_clip_evidence_dossier(clip_obs) damage_map = {"none": 0.0, "minor": 0.05, "moderate": 0.20, "severe": 0.50} marginal_damage = damage_map.get(dossier.identity_drift_severity, 0.1) clip_disp_obs = ClipDispositionObservation( evidence_dossier=dossier, minimum_clips_needed=20, phoneme_gap_severity={}, pose_gap_severity={}, budget_remaining=max(0, 50 - clip_obs.clips_audited_so_far), marginal_training_damage=marginal_damage, marginal_coverage_gain=float(clip_obs.phoneme_coverage_new), ) clip_action = recommend_clip_disposition(clip_disp_obs) clip_gt = self._synthesize_clip_ground_truth(clip_action) scores.append(grade_clip_disposition(clip_action, clip_gt)) except Exception as exc: # noqa: BLE001 log.warning("Sub-env 2 custom clip evaluation failed: %s", exc) return sum(scores) / len(scores) if scores else 0.5 def _run_subenv2_internally(self) -> float: if self._custom_clip_obs is not None: return self._run_subenv2_from_observations(self._custom_clip_obs) if not self._subenv2_cases: return 0.5 scores: list[float] = [] for clip_case in self._subenv2_cases[:5]: try: clip_obs = ClipSignalObservation.model_validate(clip_case["observation"]) clip_gt = GroundTruthClipAnnotation.model_validate(clip_case["ground_truth"]) dossier = self._build_clip_evidence_dossier(clip_obs) damage_map = {"none": 0.0, "minor": 0.05, "moderate": 0.20, "severe": 0.50} marginal_damage = damage_map.get(dossier.identity_drift_severity, 0.1) clip_disp_obs = ClipDispositionObservation( evidence_dossier=dossier, minimum_clips_needed=20, phoneme_gap_severity={}, pose_gap_severity={}, budget_remaining=max(0, 50 - clip_obs.clips_audited_so_far), marginal_training_damage=marginal_damage, marginal_coverage_gain=float(clip_obs.phoneme_coverage_new), ) clip_action = recommend_clip_disposition(clip_disp_obs) scores.append(grade_clip_disposition(clip_action, clip_gt)) except Exception as exc: # noqa: BLE001 log.warning("Sub-env 2 internal clip evaluation failed: %s", exc) return sum(scores) / len(scores) if scores else 0.5 def _build_weight_evidence_from_observation( self, weight_obs: WeightSignalObservation, ) -> WeightEvidenceDossier: rank_values = list(weight_obs.layer_rank_utilization.values()) mean_rank_util = sum(rank_values) / max(1, len(rank_values)) if weight_obs.overfitting_signature >= 0.6: training_quality = "overfit" elif weight_obs.gradient_noise_estimate >= 0.5: training_quality = "unstable" elif mean_rank_util <= 0.3: training_quality = "underfit" else: training_quality = "healthy" if mean_rank_util >= 0.75: rank_assessment = "efficient" elif mean_rank_util >= 0.45: rank_assessment = "wasteful" else: rank_assessment = "collapsed" token_map = weight_obs.token_position_to_phoneme or {} high_entropy_flags: list[TokenAnomalyFlag] = [] for token_pos in weight_obs.high_entropy_token_positions[:64]: severity = min(1.0, 0.65 + 0.05 * (token_pos % 5)) high_entropy_flags.append( TokenAnomalyFlag( token_position=int(token_pos), mapped_phoneme=token_map.get(int(token_pos)), anomaly_type="excessive_influence", severity=severity, evidence="High canonical directional anomaly score for token position.", ) ) layer_flags: list[LayerAnomalyFlag] = [] for layer_name, util in weight_obs.layer_rank_utilization.items(): sparsity = float(weight_obs.layer_sparsity.get(layer_name, 0.0)) layer_norm = float(weight_obs.layer_norms.get(layer_name, 0.0)) if util < 0.35: layer_flags.append( LayerAnomalyFlag( layer_name=layer_name, anomaly_type="rank_collapse", severity=min(1.0, float(1.0 - util)), evidence=f"Rank utilization dropped to {util:.3f}.", ) ) if sparsity > 0.85: layer_flags.append( LayerAnomalyFlag( layer_name=layer_name, anomaly_type="sparsity_anomaly", severity=min(1.0, sparsity), evidence=f"Layer sparsity elevated at {sparsity:.3f}.", ) ) if layer_norm > 8.0: layer_flags.append( LayerAnomalyFlag( layer_name=layer_name, anomaly_type="norm_explosion", severity=min(1.0, layer_norm / 12.0), evidence=f"Layer norm elevated at {layer_norm:.3f}.", ) ) global_risk = max( float(weight_obs.overfitting_signature), float(weight_obs.gradient_noise_estimate), max((flag.severity for flag in high_entropy_flags), default=0.0), max((flag.severity for flag in layer_flags), default=0.0), ) if global_risk < 0.30: overall_behavioral_risk = "low" elif global_risk < 0.55: overall_behavioral_risk = "medium" elif global_risk < 0.80: overall_behavioral_risk = "high" else: overall_behavioral_risk = "critical" evidence_summary = ( f"training_quality={training_quality}, mean_rank_util={mean_rank_util:.3f}, " f"noise={weight_obs.gradient_noise_estimate:.3f}, overfit={weight_obs.overfitting_signature:.3f}." ) return WeightEvidenceDossier( weight_file_id=weight_obs.weight_file_id, training_quality=training_quality, rank_utilization_assessment=rank_assessment, high_entropy_token_flags=high_entropy_flags, layer_anomaly_flags=layer_flags, overall_behavioral_risk=overall_behavioral_risk, evidence_summary=evidence_summary, ) def _build_phoneme_risk_observation_from_weight( self, weight_obs: WeightSignalObservation, ) -> PhonemeRiskObservation: weight_evidence = self._build_weight_evidence_from_observation(weight_obs) token_map = weight_obs.token_position_to_phoneme or {} phoneme_to_token_indices: dict[str, list[int]] = {} for token_idx, phoneme in token_map.items(): phoneme_to_token_indices.setdefault(str(phoneme), []).append(int(token_idx)) phoneme_vocabulary = sorted(phoneme_to_token_indices.keys()) if not phoneme_vocabulary: phoneme_vocabulary = [ "AA", "AE", "AH", "AO", "AW", "AY", "EH", "ER", "EY", "IH", "IY", "OW", "OY", "UH", "UW", ] phoneme_to_token_indices = {ph: [] for ph in phoneme_vocabulary} high_entropy_set = set(weight_obs.high_entropy_token_positions) phoneme_entropy_scores: dict[str, float] = {} phoneme_influence_scores: dict[str, float] = {} for phoneme in phoneme_vocabulary: indices = phoneme_to_token_indices.get(phoneme, []) flagged = sum(1 for idx in indices if idx in high_entropy_set) ratio = flagged / max(1, len(indices)) entropy_score = min( 1.0, 0.25 + 0.55 * ratio + 0.20 * float(weight_obs.overfitting_signature), ) influence_score = min( 1.0, 0.20 + 0.50 * ratio + 0.30 * float(weight_obs.gradient_noise_estimate), ) phoneme_entropy_scores[phoneme] = round(float(entropy_score), 4) phoneme_influence_scores[phoneme] = round(float(influence_score), 4) sorted_risk = sorted( phoneme_vocabulary, key=lambda ph: phoneme_entropy_scores[ph] + phoneme_influence_scores[ph], reverse=True, ) cooccurrence: list[tuple[str, str, float]] = [] for idx in range(max(0, min(3, len(sorted_risk) - 1))): left = sorted_risk[idx] right = sorted_risk[idx + 1] pair_score = ( phoneme_entropy_scores[left] + phoneme_influence_scores[left] + phoneme_entropy_scores[right] + phoneme_influence_scores[right] ) / 4.0 if pair_score > 0.60: cooccurrence.append((left, right, round(float(pair_score), 4))) training_distribution: dict[str, int] | None = None if self._custom_clip_obs: dist: dict[str, int] = {} for clip_obs in self._custom_clip_obs: for phoneme in clip_obs.phoneme_sequence: dist[phoneme] = int(dist.get(phoneme, 0)) + 1 training_distribution = dist if dist else None return PhonemeRiskObservation( weight_evidence=weight_evidence, high_entropy_token_flags=weight_evidence.high_entropy_token_flags, phoneme_vocabulary=phoneme_vocabulary, phoneme_to_token_indices=phoneme_to_token_indices, phoneme_entropy_scores=phoneme_entropy_scores, phoneme_influence_scores=phoneme_influence_scores, phoneme_cooccurrence_anomalies=cooccurrence, behavior_vocabulary=["smile", "blink", "head_turn", "jaw_drift", "brow_raise"], training_data_phoneme_distribution=training_distribution, suspected_anomalous_phonemes_from_subenv2=( weight_obs.suspected_anomalous_phonemes if weight_obs.suspected_anomalous_phonemes else None ), ) def _synthesize_behavioral_ground_truth( self, baseline_action: PhonemeRiskAction, ) -> GroundTruthBehavioralAnnotation: return GroundTruthBehavioralAnnotation( phoneme_risk_ranking=baseline_action.phoneme_risk_ranking, predicted_behavior_triggers=baseline_action.predicted_behavior_triggers, risky_phoneme_clusters=baseline_action.risky_phoneme_clusters, model_behavioral_safety=baseline_action.model_behavioral_safety, valid_mitigation_set={ (rec.target, rec.action) for rec in baseline_action.mitigation_recommendations }, ) def _build_phoneme_risk_observation(self) -> PhonemeRiskObservation: if self._custom_weight_obs is not None: return self._build_phoneme_risk_observation_from_weight(self._custom_weight_obs) if self._subenv3_case and isinstance(self._subenv3_case.get("observation"), dict): return PhonemeRiskObservation.model_validate(self._subenv3_case["observation"]) if self._subenv3_case and isinstance(self._subenv3_case.get("phoneme_obs"), dict): return PhonemeRiskObservation.model_validate(self._subenv3_case["phoneme_obs"]) return PhonemeRiskObservation.model_validate( { "weight_evidence": { "weight_file_id": "fallback.safetensors", "training_quality": "healthy", "rank_utilization_assessment": "efficient", "high_entropy_token_flags": [], "layer_anomaly_flags": [], "overall_behavioral_risk": "low", "evidence_summary": "No significant anomalies detected.", }, "high_entropy_token_flags": [], "phoneme_vocabulary": [ "AA", "AE", "AH", "AO", "AW", "AY", "EH", "ER", "EY", "IH", "IY", "OW", "OY", "UH", "UW", ], "phoneme_to_token_indices": {}, "phoneme_entropy_scores": {}, "phoneme_influence_scores": {}, "phoneme_cooccurrence_anomalies": [], "behavior_vocabulary": ["smile", "blink", "head_turn"], "training_data_phoneme_distribution": None, "suspected_anomalous_phonemes_from_subenv2": None, } ) @property def state(self) -> State: return self._state