Spaces:
Sleeping
Sleeping
| """ | |
| 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. | |
| Final reward: 0.25 * subenv1 + 0.35 * subenv2 + 0.40 * subenv3 | |
| """ | |
| from __future__ import annotations | |
| from typing import 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 | |
| """ | |
| pass | |