Type EpisodeState and EpisodeLog with Protocol, ConversationEntry, RewardBreakdown (MOD 04)

Replace loose dict fields with typed Pydantic models:
- EpisodeState.current_protocol: Optional[Protocol]
- EpisodeState.conversation_history: list[ConversationEntry]
- EpisodeLog.transcript: list[ConversationEntry]
- EpisodeLog.reward_breakdown: Optional[RewardBreakdown]

Stub server now explicitly constructs Protocol and ConversationEntry
objects instead of raw dicts. Fix WS handler info serialization.
Add 8 tests covering typed construction, JSON round-trip, and nesting.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

replicalab/models.py

@@ -374,8 +374,8 @@ class EpisodeState(BaseModel):
     lab_reagents: list[str] = Field(default_factory=list)
     lab_staff_count: int = 0
     lab_time_limit_days: int = 0
-    current_protocol: Optional[dict] = None
-    conversation_history: list[dict] = Field(default_factory=list)
+    current_protocol: Optional[Protocol] = None
+    conversation_history: list[ConversationEntry] = Field(default_factory=list)
     round_number: int = 0
     max_rounds: int = 0
     done: bool = False
@@ -394,8 +394,8 @@ class EpisodeLog(BaseModel):
     scenario_template: str = ""
     difficulty: str = "easy"
     final_state: Optional[EpisodeState] = None
-    transcript: list[dict] = Field(default_factory=list)
-    reward_breakdown: dict = Field(default_factory=dict)
+    transcript: list[ConversationEntry] = Field(default_factory=list)
+    reward_breakdown: Optional[RewardBreakdown] = None
     total_reward: float = 0.0
     rounds_used: int = 0
     agreement_reached: bool = False

server/app.py

@@ -42,10 +42,12 @@ from replicalab.config import (
 )
 from replicalab.scenarios import available_scenario_families, generate_scenario
 from replicalab.models import (
+    ConversationEntry,
     EpisodeLog,
     EpisodeState,
     LabManagerObservation,
     Observation,
+    Protocol,
     RewardBreakdown,
     ScientistAction,
     ScientistObservation,
@@ -117,7 +119,7 @@ class _StubEnv:
 
     def __init__(self) -> None:
         self._state = EpisodeState()
-        self._logs: list[dict] = []
+        self._logs: list[ConversationEntry] = []
         self._episode_id: str = ""
 
     # ── public interface (matches ReplicaLabEnv) ──────────────────────────
@@ -202,46 +204,46 @@ class _StubEnv:
 
     # ── internal helpers ──────────────────────────────────────────────────
 
-    def _scientist_log_entry(self, action: ScientistAction) -> dict[str, Any]:
+    def _scientist_log_entry(self, action: ScientistAction) -> ConversationEntry:
         action_type = (
             action.action_type.value
             if hasattr(action.action_type, "value")
             else str(action.action_type)
         )
         message = action.rationale or f"Scientist chose action '{action_type}'."
-        return {
-            "role": "scientist",
-            "message": message,
-            "round_number": self._state.round_number,
-            "action_type": action_type,
-        }
-
-    def _lab_manager_log_entry(self, action: ScientistAction) -> dict[str, Any]:
+        return ConversationEntry(
+            role="scientist",
+            message=message,
+            round_number=self._state.round_number,
+            action_type=action_type,
+        )
+
+    def _lab_manager_log_entry(self, action: ScientistAction) -> ConversationEntry:
         if action.action_type == "accept":
             message = "Stub review: agreement recorded and episode will close."
             action_type = "accept"
         else:
             message = "Stub review: proposal received and remains feasible under the stub lab."
             action_type = "report_feasibility"
-        return {
-            "role": "lab_manager",
-            "message": message,
-            "round_number": self._state.round_number,
-            "action_type": action_type,
-        }
-
-    def _protocol_from_action(self, action: ScientistAction) -> dict[str, Any] | None:
+        return ConversationEntry(
+            role="lab_manager",
+            message=message,
+            round_number=self._state.round_number,
+            action_type=action_type,
+        )
+
+    def _protocol_from_action(self, action: ScientistAction) -> Optional[Protocol]:
         if action.action_type not in {"propose_protocol", "revise_protocol"}:
             return self._state.current_protocol
-        return {
-            "technique": action.technique,
-            "sample_size": action.sample_size,
-            "controls": list(action.controls),
-            "duration_days": action.duration_days,
-            "required_equipment": list(action.required_equipment),
-            "required_reagents": list(action.required_reagents),
-            "rationale": action.rationale,
-        }
+        return Protocol(
+            technique=action.technique,
+            sample_size=action.sample_size,
+            controls=list(action.controls),
+            duration_days=action.duration_days,
+            required_equipment=list(action.required_equipment),
+            required_reagents=list(action.required_reagents),
+            rationale=action.rationale,
+        )
 
     def _make_observation(self) -> Observation:
         s = self._state
@@ -572,7 +574,7 @@ async def websocket_endpoint(ws: WebSocket):
                             else None,
                             "reward": result.reward,
                             "done": result.done,
-                            "info": result.info,
+                            "info": result.info.model_dump(),
                         },
                     )
                 except Exception as exc:

tests/test_models.py

@@ -5,14 +5,19 @@ from pydantic import ValidationError
 
 from replicalab.models import (
     ConversationEntry,
+    EpisodeLog,
+    EpisodeState,
     LabManagerAction,
     LabManagerActionType,
     LabManagerObservation,
     Observation,
     Protocol,
+    RewardBreakdown,
     ScientistAction,
     ScientistActionType,
     ScientistObservation,
+    StepInfo,
+    StepResult,
 )
 
 
@@ -230,3 +235,178 @@ def test_observation_rejects_negative_budget() -> None:
 
     with pytest.raises(ValidationError):
         Observation.model_validate(payload)
+
+
+# ---------------------------------------------------------------------------
+# MOD 04 — Typed EpisodeState and EpisodeLog
+# ---------------------------------------------------------------------------
+
+
+def _sample_protocol() -> Protocol:
+    return Protocol(
+        sample_size=32,
+        controls=["vehicle_control", "positive_control"],
+        technique="manual_cell_counting",
+        duration_days=5,
+        required_equipment=["microscope", "co2_incubator"],
+        required_reagents=["dmso", "drug_x", "culture_media"],
+        rationale="Uses available equipment while preserving controls.",
+    )
+
+
+def _sample_conversation_entry() -> ConversationEntry:
+    return ConversationEntry(
+        role="scientist",
+        message="I propose a manual counting protocol.",
+        round_number=1,
+        action_type="propose_protocol",
+    )
+
+
+def test_episode_state_accepts_typed_protocol_and_history() -> None:
+    protocol = _sample_protocol()
+    entry = _sample_conversation_entry()
+    state = EpisodeState(
+        seed=42,
+        current_protocol=protocol,
+        conversation_history=[entry],
+        round_number=1,
+        max_rounds=6,
+    )
+
+    assert isinstance(state.current_protocol, Protocol)
+    assert state.current_protocol.technique == "manual_cell_counting"
+    assert isinstance(state.conversation_history[0], ConversationEntry)
+    assert state.conversation_history[0].role == "scientist"
+
+
+def test_episode_state_accepts_none_protocol() -> None:
+    state = EpisodeState(current_protocol=None, conversation_history=[])
+    assert state.current_protocol is None
+    assert state.conversation_history == []
+
+
+def test_episode_state_json_round_trip() -> None:
+    protocol = _sample_protocol()
+    entry = _sample_conversation_entry()
+    state = EpisodeState(
+        seed=7,
+        scenario_template="math_reasoning",
+        difficulty="hard",
+        paper_title="Test Paper",
+        current_protocol=protocol,
+        conversation_history=[entry],
+        round_number=2,
+        max_rounds=6,
+    )
+
+    dumped = state.model_dump_json()
+    restored = EpisodeState.model_validate_json(dumped)
+
+    assert isinstance(restored.current_protocol, Protocol)
+    assert restored.current_protocol.sample_size == 32
+    assert isinstance(restored.conversation_history[0], ConversationEntry)
+    assert restored.conversation_history[0].action_type == "propose_protocol"
+    assert restored.seed == 7
+
+
+def test_episode_log_accepts_typed_fields() -> None:
+    entry = _sample_conversation_entry()
+    breakdown = RewardBreakdown(rigor=0.8, feasibility=0.7, fidelity=0.9)
+    log = EpisodeLog(
+        episode_id="ep-001",
+        seed=42,
+        transcript=[entry],
+        reward_breakdown=breakdown,
+        total_reward=5.0,
+        rounds_used=3,
+        agreement_reached=True,
+    )
+
+    assert isinstance(log.transcript[0], ConversationEntry)
+    assert isinstance(log.reward_breakdown, RewardBreakdown)
+    assert log.reward_breakdown.rigor == 0.8
+
+
+def test_episode_log_none_reward_breakdown() -> None:
+    log = EpisodeLog(episode_id="ep-002")
+    assert log.reward_breakdown is None
+    assert log.transcript == []
+
+
+def test_episode_log_json_round_trip() -> None:
+    entry = _sample_conversation_entry()
+    breakdown = RewardBreakdown(
+        rigor=0.6, feasibility=0.5, fidelity=0.7,
+        efficiency_bonus=0.1, communication_bonus=0.05,
+        penalties={"timeout": 0.02},
+    )
+    state = EpisodeState(
+        seed=99,
+        current_protocol=_sample_protocol(),
+        conversation_history=[entry],
+        round_number=3,
+        max_rounds=6,
+        done=True,
+        agreement_reached=True,
+        reward=5.0,
+        rigor_score=0.6,
+    )
+    log = EpisodeLog(
+        episode_id="ep-round-trip",
+        seed=99,
+        final_state=state,
+        transcript=[entry],
+        reward_breakdown=breakdown,
+        total_reward=5.0,
+        rounds_used=3,
+        agreement_reached=True,
+        judge_notes="Good protocol.",
+        verdict="accept",
+    )
+
+    dumped = log.model_dump_json()
+    restored = EpisodeLog.model_validate_json(dumped)
+
+    assert isinstance(restored.final_state, EpisodeState)
+    assert isinstance(restored.final_state.current_protocol, Protocol)
+    assert isinstance(restored.final_state.conversation_history[0], ConversationEntry)
+    assert isinstance(restored.transcript[0], ConversationEntry)
+    assert isinstance(restored.reward_breakdown, RewardBreakdown)
+    assert restored.reward_breakdown.penalties == {"timeout": 0.02}
+    assert restored.episode_id == "ep-round-trip"
+
+
+def test_episode_log_nested_state_preserves_typed_fields() -> None:
+    protocol = _sample_protocol()
+    entry = _sample_conversation_entry()
+    state = EpisodeState(
+        current_protocol=protocol,
+        conversation_history=[entry],
+    )
+    log = EpisodeLog(final_state=state)
+
+    assert isinstance(log.final_state.current_protocol, Protocol)
+    assert log.final_state.current_protocol.technique == "manual_cell_counting"
+    assert isinstance(log.final_state.conversation_history[0], ConversationEntry)
+
+
+def test_step_result_with_typed_info() -> None:
+    breakdown = RewardBreakdown(rigor=0.8, feasibility=0.8, fidelity=0.8)
+    info = StepInfo(
+        agreement_reached=True,
+        reward_breakdown=breakdown,
+        judge_notes="All checks passed.",
+        verdict="accept",
+        round=3,
+        stub=True,
+    )
+    result = StepResult(reward=5.0, done=True, info=info)
+
+    dumped = result.model_dump_json()
+    restored = StepResult.model_validate_json(dumped)
+
+    assert isinstance(restored.info, StepInfo)
+    assert isinstance(restored.info.reward_breakdown, RewardBreakdown)
+    assert restored.info.agreement_reached is True
+    assert restored.info.verdict == "accept"