""" Data models for the TalkingHeadBench environment. TalkingHeadBench evaluates diagnostic reasoning across 3 coupled sub-environments for talking-head LoRA pipelines. Each episode has three agent-facing decision points: Step 1 — Image Diagnostician (Node 1): ImageDiagnosticsAction Step 2 — Param Anomaly Detector (Node 2): ParamAnomalyAction Step 3 — Phoneme Risk Assessor (Node 8): PhonemeRiskAction All grading is deterministic and rule-based. No LLM judge required. Full-pipeline final reward: 0.25 * subenv1 + 0.35 * subenv2 + 0.40 * subenv3. Mode-based OpenEnv episodes return mode-local rewards (or a 50/50 subenv2/subenv3 blend for clips_and_weights). """ from __future__ import annotations from typing import Any, Literal, Optional from openenv.core.env_server.types import Action, Observation from pydantic import Field # --------------------------------------------------------------------------- # Shared sub-types (inlined to keep models.py self-contained) # --------------------------------------------------------------------------- # --------------------------------------------------------------------------- # Step 1 — Image Diagnostician Action (Node 1) # --------------------------------------------------------------------------- class ImageDiagnosticsAction(Action): """ Structured diagnosis of a reference image and prompt produced by the Image Diagnostician agent (Node 1). """ regime_classification: Literal[ "frontal_simple", "non_frontal", "complex_background", "occluded", "low_quality", ] = Field(..., description="Regime classification of the reference image.") identified_risk_factors: list[str] = Field( ..., description="Specific image or prompt risk factors detected." ) prompt_issues: list[str] = Field( default_factory=list, description="Conflicting or weakly-anchored terms in the prompt.", ) recommended_prompt_modifications: list[str] = Field( default_factory=list, description="Actionable suggestions to improve the prompt.", ) image_usability_score: float = Field( ..., ge=0.0, le=1.0, description="Overall image usability in [0, 1]." ) reasoning: str = Field( ..., description="Brief rationale for the regime and primary risk factor." ) # --------------------------------------------------------------------------- # Step 2 — Parameter Anomaly Detector Action (Node 2) # --------------------------------------------------------------------------- class ParameterAnomaly(Action): """A single flagged parameter anomaly.""" parameter: str = Field(..., description='Parameter name, e.g. "cfg" or "eta".') issue: str = Field(..., description="Description of the anomaly.") severity: Literal["minor", "moderate", "severe"] = Field(..., description="Severity level.") linked_failure_mode: str = Field(..., description="Associated generation failure mode.") class DirectionalFix(Action): """A directional (non-prescriptive) fix recommendation.""" target: str = Field(..., description="Parameter to fix.") direction: Literal[ "increase", "decrease", "enable", "disable", "reconsider" ] = Field(..., description="Direction of the fix.") rationale: str = Field(..., description="Why this fix is recommended.") priority: Literal["critical", "recommended", "optional"] = Field( ..., description="Fix priority." ) class ParamAnomalyAction(Action): """ Structured output of the Parameter Anomaly Detector agent (Node 2). """ config_risk_level: Literal["safe", "marginal", "risky", "dangerous"] = Field( ..., description="Overall risk level of the proposed configuration." ) anomalies: list[ParameterAnomaly] = Field( default_factory=list, description="Flagged parameter anomalies." ) predicted_failure_modes: list[ Literal[ "identity_collapse", "reference_token_dropout", "temporal_jitter", "background_bleed", "lip_sync_desync", "pose_instability", "overexposure_artifacts", ] ] = Field(default_factory=list, description="Predicted generation failure modes.") directional_fixes: list[DirectionalFix] = Field( default_factory=list, description="Directional fix recommendations." ) summary: str = Field(..., description="One-sentence risk summary.") # --------------------------------------------------------------------------- # Step 3 — Phoneme Risk Assessor Action (Node 8) # --------------------------------------------------------------------------- class PhonemeRiskEntry(Action): """Risk profile for a single phoneme.""" phoneme: str = Field(..., description="Phoneme string, e.g. 'AH' or 'S'.") risk_score: float = Field(..., ge=0.0, le=1.0, description="Risk score in [0, 1].") risk_type: Literal[ "identity_trigger", "expression_trigger", "motion_trigger", "artifact_trigger", "unknown_anomaly", ] = Field(..., description="Category of behavioral risk.") confidence: float = Field(..., ge=0.0, le=1.0, description="Confidence in [0, 1].") evidence: str = Field(..., description="Evidence justifying the risk score.") class BehaviorTriggerPrediction(Action): """Predicted phoneme → behavior association.""" trigger_phoneme: str = Field(..., description="Phoneme that triggers the behavior.") triggered_behavior: str = Field( ..., description='Behavior label, e.g. "smile" or "blink".' ) association_strength: float = Field(..., ge=0.0, le=1.0) is_intended: bool = Field(..., description="Whether the association is intentional.") concern_level: Literal["none", "low", "medium", "high"] = Field(..., description="Concern level.") class PhonemeCluster(Action): """A cluster of phonemes sharing a common risk pattern.""" phonemes: list[str] = Field(..., description="Phonemes in the cluster.") cluster_risk_type: str = Field(..., description="Description of the shared risk pattern.") combined_risk_score: float = Field(..., ge=0.0, le=1.0) interaction_description: str = Field(..., description="How the phonemes interact.") class MitigationRecommendation(Action): """A directional mitigation for a phoneme-level behavioral risk.""" target: str = Field(..., description="Phoneme or layer to target.") action: Literal[ "retrain_with_more_data", "remove_from_dataset", "add_counter_examples", "reduce_lora_rank", "apply_weight_regularization", "flag_for_manual_review", ] = Field(..., description="Recommended mitigation action.") rationale: str = Field(..., description="Why this mitigation is recommended.") priority: Literal["critical", "recommended", "optional"] = Field(..., description="Priority level.") class PhonemeRiskAction(Action): """ Structured behavioral risk profile produced by the Phoneme Risk Assessor (Node 8). """ phoneme_risk_ranking: list[PhonemeRiskEntry] = Field( default_factory=list, description="Ranked list of at-risk phonemes." ) predicted_behavior_triggers: list[BehaviorTriggerPrediction] = Field( default_factory=list, description="Predicted phoneme-to-behavior associations." ) risky_phoneme_clusters: list[PhonemeCluster] = Field( default_factory=list, description="Clusters of phonemes with shared risk patterns." ) model_behavioral_safety: Literal[ "safe", "minor_concerns", "moderate_risk", "high_risk", "unsafe" ] = Field(..., description="Overall behavioral safety rating.") mitigation_recommendations: list[MitigationRecommendation] = Field( default_factory=list, description="Recommended mitigations." ) summary: str = Field(..., description="One-paragraph summary of behavioral risks.") # --------------------------------------------------------------------------- # Shared Observation type # --------------------------------------------------------------------------- class TalkingHeadObservation(Observation): """ Observation emitted by TalkingHeadBench at each step. The ``metadata`` dict carries the full observation payload: - ``node``: which node produced this observation - ``observation``: the typed signal dict for the agent - ``expected_action_schema``: name of the Action class to return - ``instruction``: natural-language task description - ``step``: current episode step index (0–3) - ``scores`` (on done=True): per-sub-env and final score breakdown """ node: str = Field(default="", description="Node identifier for the current step.") step: int = Field(default=0, description="Episode step index (0-based).") api_version: str = Field(default="1.0", description="Observation schema version.") mode: Literal["benchmark", "custom"] = Field( default="benchmark", description="Episode mode: benchmark test-set execution or custom ingested bundle.", ) is_deterministic: bool = Field( default=True, description="Whether the episode scoring behavior is deterministic for fixed inputs.", ) case_id: Optional[str] = Field(default=None, description="Case identifier for the episode.") episode_id: Optional[str] = Field(default=None, description="Episode session identifier.") instruction: Optional[str] = Field( default=None, description="Natural-language instruction for the current node decision.", ) expected_action_schema: Optional[str] = Field( default=None, description="Name of the expected action schema for this step.", ) signals: dict[str, Any] = Field( default_factory=dict, description="Agent-facing pre-extracted signal dictionary.", ) provenance: dict[str, Any] = Field( default_factory=dict, description="Source and provenance metadata for observations and ground truth.", ) scores: Optional[dict[str, Any]] = Field( default=None, description="Score breakdown populated at episode completion.", ) reward_formula: Optional[str] = Field( default=None, description="Reward formula description populated on done=True.", ) error: Optional[str] = Field(default=None, description="Optional terminal error details.")