Spaces:
Sleeping
Sleeping
Shivappa Sangapur
feat: DebateArenaEnv complete - server/client separation, LLM multi-agent, Docker, OpenEnv compliant
b1eddd6 | """ | |
| DebateArenaEnv — Core RL environment. | |
| Implements the OpenEnv gym-style contract: | |
| env.reset(topic_id) → observation | |
| env.step(action) → observation, reward, done, info | |
| env.state() → current state dict | |
| env.close() → None | |
| AgentRole : PROPOSER | CHALLENGER | |
| EpisodePhase: OPEN | DONE | |
| """ | |
| from __future__ import annotations | |
| import time | |
| from dataclasses import dataclass, field | |
| from enum import Enum | |
| from typing import Any, Dict, List, Optional, Tuple | |
| from server.tasks import REGISTRY, TOPIC_BANK, Topic | |
| from server.rubric import DebateArenaRubric | |
| # ── OpenEnv base shim ───────────────────────────────────────────────────────── | |
| try: | |
| from openenv import Environment as _OpenEnvBase # type: ignore | |
| except ImportError: | |
| class _OpenEnvBase: # type: ignore | |
| """Minimal shim when openenv package is not installed.""" | |
| pass | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| # Enums | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| class AgentRole(str, Enum): | |
| PROPOSER = "proposer" | |
| CHALLENGER = "challenger" | |
| class EpisodePhase(str, Enum): | |
| OPEN = "open" | |
| DONE = "done" | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| # Tool permission sets | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| PROPOSER_TOOLS = { | |
| "get_topic", "fact_check", | |
| "make_argument", "update_position", | |
| "concede_sub_point", "close_debate", | |
| } | |
| CHALLENGER_TOOLS = { | |
| "get_topic", "fact_check", | |
| "challenge_argument", "close_debate", | |
| } | |
| # Tools that do NOT advance the round counter | |
| INFO_TOOLS = {"get_topic", "fact_check"} | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| # Dataclasses | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| class AgentAction: | |
| role: AgentRole | |
| tool: str | |
| params: Dict[str, Any] | |
| timestamp: float = field(default_factory=time.time) | |
| class EpisodeState: | |
| topic: Topic | |
| phase: EpisodePhase = EpisodePhase.OPEN | |
| round_number: int = 0 | |
| proposer_args: List[str] = field(default_factory=list) | |
| challenger_args: List[str] = field(default_factory=list) | |
| facts_cited: List[str] = field(default_factory=list) | |
| position_updated: bool = False | |
| conceded: bool = False | |
| winner: Optional[str] = None | |
| final_reward: float = 0.0 | |
| history: List[AgentAction] = field(default_factory=list) | |
| def debate_closed(self) -> bool: | |
| return self.phase == EpisodePhase.DONE | |
| def all_text(self) -> str: | |
| return " ".join(self.proposer_args + self.challenger_args) | |
| def to_dict(self) -> Dict[str, Any]: | |
| return { | |
| "phase": self.phase.value, | |
| "round_number": self.round_number, | |
| "proposer_args": self.proposer_args, | |
| "challenger_args": self.challenger_args, | |
| "facts_cited": self.facts_cited, | |
| "position_updated": self.position_updated, | |
| "conceded": self.conceded, | |
| "winner": self.winner, | |
| "final_reward": self.final_reward, | |
| "all_text": self.all_text(), | |
| "rounds_used": self.round_number, | |
| } | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| # Environment | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| class DebateArenaEnv(_OpenEnvBase): | |
| """ | |
| Multi-agent epistemic debate environment. | |
| Two agents (Proposer, Challenger) take turns using MCP tools. | |
| The episode ends when either agent calls end_debate or max_rounds is hit. | |
| Reward is computed by DebateArenaRubric at episode close. | |
| """ | |
| def __init__(self) -> None: | |
| self._rubric = DebateArenaRubric() | |
| self._state: Optional[EpisodeState] = None | |
| self._topic: Optional[Topic] = None | |
| # ── Public API ──────────────────────────────────────────────────────────── | |
| def reset(self, topic_id: Optional[str] = None) -> Dict[str, Any]: | |
| """Start a new episode. Returns initial observation.""" | |
| if topic_id: | |
| topic = REGISTRY.get_topic(topic_id) | |
| if topic is None: | |
| topic = REGISTRY.current_topic() | |
| else: | |
| topic = REGISTRY.current_topic() | |
| self._topic = topic | |
| self._state = EpisodeState(topic=topic) | |
| return self._observation("Episode started. Proposer goes first.") | |
| def step(self, action: AgentAction) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]: | |
| """ | |
| Execute one agent action. | |
| Returns (observation, reward, done, info). | |
| """ | |
| if self._state is None: | |
| raise RuntimeError("Call reset() before step().") | |
| if self._state.debate_closed: | |
| raise RuntimeError("Episode is already done. Call reset().") | |
| self._validate_tool(action) | |
| self._state.history.append(action) | |
| # Dispatch to internal handler | |
| result_msg, done = self._dispatch(action) | |
| # Advance round counter only for non-info tools | |
| if action.tool not in INFO_TOOLS: | |
| self._state.round_number += 1 | |
| # Auto-close on max_rounds | |
| if self._state.round_number >= self._topic.max_rounds and not done: # type: ignore[union-attr] | |
| done = True | |
| result_msg += " (Max rounds reached — debate auto-closed.)" | |
| reward = 0.01 # interim reward; final computed on close | |
| if done: | |
| reward, _ = self._close_episode() | |
| obs = self._observation(result_msg) | |
| return obs, reward, done, {"result": result_msg} | |
| def state(self) -> Dict[str, Any]: | |
| if self._state is None: | |
| return {"error": "No active episode. Call reset() first."} | |
| return { | |
| "topic": self._topic.to_dict() if self._topic else {}, # type: ignore[union-attr] | |
| "state": self._state.to_dict(), | |
| } | |
| def close(self) -> None: | |
| self._state = None | |
| self._topic = None | |
| def render(self) -> str: | |
| if self._state is None: | |
| return "No active episode." | |
| s = self._state | |
| lines = [ | |
| f"=== DebateArenaEnv ===", | |
| f"Topic : {self._topic.claim if self._topic else '—'}", # type: ignore[union-attr] | |
| f"Phase : {s.phase.value} Round: {s.round_number}", | |
| f"Proposer args : {len(s.proposer_args)}", | |
| f"Challenger args: {len(s.challenger_args)}", | |
| f"Facts cited : {s.facts_cited}", | |
| f"Updated position: {s.position_updated} Conceded: {s.conceded}", | |
| ] | |
| return "\n".join(lines) | |
| # ── Internal helpers ────────────────────────────────────────────────────── | |
| def _validate_tool(self, action: AgentAction) -> None: | |
| """Raise ValueError if the tool is not permitted for this role.""" | |
| # We allow all tools via HTTP dispatch; role enforcement is advisory here | |
| pass | |
| def _dispatch(self, action: AgentAction) -> Tuple[str, bool]: | |
| """Route action.tool → handler. Returns (message, done).""" | |
| tool = action.tool | |
| params = action.params | |
| state = self._state | |
| if tool == "verify_fact": | |
| fact_key = params.get("fact_key", "") | |
| known = self._topic.known_facts if self._topic else {} # type: ignore[union-attr] | |
| if fact_key in known: | |
| verdict = "✅ TRUE" if known[fact_key] else "❌ FALSE" | |
| safe = "Safe to cite." if known[fact_key] else "Do NOT cite — will trigger hallucination penalty!" | |
| return f"[Fact Check] '{fact_key}' is {verdict}. {safe}", False | |
| return f"[Fact Check] '{fact_key}' is ❓ UNKNOWN.", False | |
| elif tool == "submit_argument": | |
| argument = params.get("argument", "") | |
| facts_cited = params.get("facts_cited", []) | |
| state.proposer_args.append(argument) # type: ignore[union-attr] | |
| state.facts_cited.extend(facts_cited) # type: ignore[union-attr] | |
| return ( | |
| f"Argument submitted (round {state.round_number + 1}). " # type: ignore[union-attr] | |
| "Challenger is preparing a rebuttal. " | |
| f"Round {state.round_number + 1}/{self._topic.max_rounds}." # type: ignore[union-attr] | |
| ), False | |
| elif tool == "submit_rebuttal": | |
| rebuttal = params.get("rebuttal", "") | |
| facts_cited = params.get("facts_cited", []) | |
| expose = params.get("expose_fallacy") | |
| state.challenger_args.append(rebuttal) # type: ignore[union-attr] | |
| state.facts_cited.extend(facts_cited) # type: ignore[union-attr] | |
| msg = "Rebuttal submitted. Proposer may now respond or update their position." | |
| if expose: | |
| msg += f" Fallacy exposed: {expose}." | |
| return msg, False | |
| elif tool == "refine_position": | |
| refined_claim = params.get("refined_claim", "") | |
| reason = params.get("reason", "") | |
| state.position_updated = True # type: ignore[union-attr] | |
| state.proposer_args.append(f"[Refined] {refined_claim} — {reason}") # type: ignore[union-attr] | |
| return "Position updated. Belief-updating reward (+0.10) will be applied at close.", False | |
| elif tool == "concede_point": | |
| sub_point = params.get("sub_point", "") | |
| maintain_main = params.get("maintain_main", "") | |
| state.conceded = True # type: ignore[union-attr] | |
| state.proposer_args.append(f"[Concede] {sub_point}. But: {maintain_main}") # type: ignore[union-attr] | |
| return "Concession noted (+0.05 concession credit). Judges reward strategic concessions.", False | |
| elif tool == "end_debate": | |
| closing = params.get("closing_statement", "") | |
| role = params.get("role", "proposer") | |
| state.proposer_args.append(f"[Closing] {closing}") # type: ignore[union-attr] | |
| state.phase = EpisodePhase.DONE # type: ignore[union-attr] | |
| return f"Debate closed by {role}.", True | |
| return f"Unknown tool: {tool}", False | |
| def _close_episode(self) -> Tuple[float, List[Any]]: | |
| """Score the episode, update curriculum, set winner.""" | |
| state = self._state | |
| topic = self._topic | |
| state.phase = EpisodePhase.DONE # type: ignore[union-attr] | |
| reward, results = self._rubric.score( | |
| topic_dict=topic.to_dict(), # type: ignore[union-attr] | |
| state_dict=state.to_dict(), # type: ignore[union-attr] | |
| ) | |
| state.final_reward = reward # type: ignore[union-attr] | |
| state.winner = "proposer" if reward >= 0.5 else "challenger" # type: ignore[union-attr] | |
| # Update adaptive curriculum | |
| REGISTRY.update_curriculum(topic.topic_id, won=(reward >= 0.5)) # type: ignore[union-attr] | |
| return reward, results | |
| def _observation(self, message: str) -> Dict[str, Any]: | |
| state = self._state | |
| topic = self._topic | |
| return { | |
| "message": message, | |
| "topic": topic.to_dict() if topic else {}, # type: ignore[union-attr] | |
| "round": state.round_number if state else 0, # type: ignore[union-attr] | |
| "done": state.debate_closed if state else False, # type: ignore[union-attr] | |
| "winner": state.winner if state else None, # type: ignore[union-attr] | |
| "reward": state.final_reward if state else 0.0, # type: ignore[union-attr] | |
| } | |
| # Singleton used by the HTTP server | |
| _ENV = DebateArenaEnv() | |
| def get_env() -> DebateArenaEnv: | |
| return _ENV | |