Spaces:
Sleeping
Sleeping
File size: 10,588 Bytes
3da2703 f755447 3da2703 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 | """
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.")
|