"""Multi-game study sessions for Prolific workers. Wraps N sequential LiveSessions (one per game) with warmup + real phases, Prolific metadata, and a completion code generated on finish. Generic across envs: each EnvRenderer supplies its own StudyConfig. """ from __future__ import annotations import random import string import threading import time from dataclasses import dataclass, field from typing import Optional from .envs import EnvRenderer, get as get_env from .sessions import LiveSession, SessionManager @dataclass class ProlificMeta: """Prolific URL params attached to a study session for trajectory metadata.""" prolific_pid: Optional[str] = None study_id: Optional[str] = None prolific_session_id: Optional[str] = None data_source: str = "" @dataclass class StudyConfig: """Per-env study setup: how many games of each kind + per-game env_kwargs.""" num_warmup: int num_real: int game_env_kwargs: list[dict] # len == num_warmup + num_real warmup_partner_keys: Optional[list[str]] = None # if None, picks randomly per game real_partner_keys: Optional[list[str]] = None # if None, picks randomly per game def total_games(self) -> int: return self.num_warmup + self.num_real @dataclass class StudySession: """One Prolific worker's run through a study. Holds a list of game session_ids, advances on completion, generates a completion code.""" study_id: str env_name: str config: StudyConfig prolific: ProlificMeta current_game_index: int = 0 game_session_ids: list[Optional[str]] = field(default_factory=list) completion_code: Optional[str] = None session_complete: bool = False last_seen_at: float = field(default_factory=time.time) def is_warmup(self) -> bool: return self.current_game_index < self.config.num_warmup def current_game_id(self) -> Optional[str]: if 0 <= self.current_game_index < len(self.game_session_ids): return self.game_session_ids[self.current_game_index] return None class StudySessionManager: def __init__(self, session_manager: SessionManager): self._sm = session_manager self._lock = threading.Lock() self._active: dict[str, StudySession] = {} def start(self, env_name: str, prolific: ProlificMeta, config: Optional[StudyConfig] = None) -> StudySession: renderer = get_env(env_name) cfg = config or renderer.study_config() if len(cfg.game_env_kwargs) != cfg.total_games(): raise ValueError( f"study config: {len(cfg.game_env_kwargs)} game_env_kwargs " f"!= {cfg.total_games()} total games" ) study_id = _gen_id("study_") study = StudySession( study_id=study_id, env_name=env_name, config=cfg, prolific=prolific, game_session_ids=[None] * cfg.total_games(), ) # Kick off game 0 immediately so the client gets a state to render. first_sid = self._launch_game(renderer, study, 0) study.game_session_ids[0] = first_sid with self._lock: self._active[study_id] = study return study def get(self, study_id: str) -> StudySession: with self._lock: s = self._active.get(study_id) if s is None: raise KeyError(f"study {study_id} expired or unknown") s.last_seen_at = time.time() return s def step(self, study_id: str, ui_action: dict) -> dict: """Step current game (mirrors /play/step's full flow: human action + partner action if their turn next). On game-done, advance to next; on study-done, generate completion code.""" study = self.get(study_id) renderer = get_env(study.env_name) cur_sid = study.current_game_id() if cur_sid is None: raise RuntimeError(f"study {study_id}: no current game") live = self._sm.get(cur_sid) from .helpers import current_player as _current_player, next_rng human_idx = renderer.HUMAN_AGENT_IDX partner_idx = 1 - human_idx # Same shape as /play/step: human action, then partner if next. action = renderer.action_from_ui(ui_action, live.state) state_before_human = live.state obs, state, reward, done, info = renderer.step( live.env, live.state, {human_idx: action}, next_rng(), ) live.state = state human_event = renderer.describe_action( action, state_before_human, live.state, float(reward), human_idx, ) events: list[dict] = [human_event] live.history.append(human_event) partner_acted = False if not done: cur_player = _current_player(live.state, human_idx) if cur_player == partner_idx: partner_action = int(live.partner.get_action(obs, live.state, next_rng())) state_before_partner = live.state obs, state, p_reward, done, info = renderer.step( live.env, live.state, {partner_idx: partner_action}, next_rng(), ) live.state = state reward += float(p_reward) partner_event = renderer.describe_action( partner_action, state_before_partner, live.state, float(p_reward), partner_idx, ) events.append(partner_event) live.history.append(partner_event) partner_acted = True advanced = False prev_idx = study.current_game_index if done: # Stamp end_time on the game that just finished (for per-game duration metric). live.end_time = time.time() if study.current_game_index < study.config.total_games() - 1: study.current_game_index += 1 next_sid = self._launch_game(renderer, study, study.current_game_index) study.game_session_ids[study.current_game_index] = next_sid advanced = True else: study.session_complete = True study.completion_code = _gen_completion_code() return { "study_id": study.study_id, "current_game_index": study.current_game_index, "total_games": study.config.total_games(), "num_warmup": study.config.num_warmup, "is_warmup": study.is_warmup(), "session_complete": study.session_complete, "completion_code": study.completion_code, "game_just_advanced": advanced, "prev_game_index": prev_idx if advanced else None, "reward": float(reward), "done": bool(done), "partner_acted": partner_acted, "events": events, } def save_all(self, study_id: str, agent_name: str = "Anonymous") -> list[dict]: """Build rich trajectory dicts for every completed game in this study. Caller (route handler) persists them via TrajectoryStore. Per-trajectory metadata mirrors Johnny's human_data_collecting save_episode: timing (start_time/end_time/duration), env_kwargs (grid_size/num_food/...), partner identity, warmup flag, sequence-aware kind tag for folder split. """ study = self.get(study_id) trajs: list[dict] = [] for idx, sid in enumerate(study.game_session_ids): if sid is None: continue try: live = self._sm.get(sid) except KeyError: continue is_warmup = idx < study.config.num_warmup env_kwargs = ( study.config.game_env_kwargs[idx] if idx < len(study.config.game_env_kwargs) else {} ) start = getattr(live, "start_time", None) end = getattr(live, "end_time", None) or time.time() traj = { "env": study.env_name, "agent_name": agent_name, "partner_key": live.partner_key, "score": live.renderer.score_summary(live.state, {}), "history": live.history, "prolific": { "prolific_pid": study.prolific.prolific_pid, "study_id": study.prolific.study_id, "prolific_session_id": study.prolific.prolific_session_id, "data_source": study.prolific.data_source, }, "study_id": study.study_id, "game_index": idx, "is_warmup": is_warmup, "kind": "warmup" if is_warmup else "real", "env_kwargs": env_kwargs, "start_time": start, "end_time": end, "duration": (end - start) if (start and end) else None, "total_steps": len(live.history), "completion_code": study.completion_code, } trajs.append(traj) return trajs def _launch_game(self, renderer: EnvRenderer, study: StudySession, idx: int) -> str: env_kwargs = study.config.game_env_kwargs[idx] partner_key = self._pick_partner(renderer, study, idx) # SessionManager.start() handles the env+partner+state setup. live, _obs = self._sm.start( env_name=study.env_name, partner_key=partner_key, env_kwargs=env_kwargs, seed=None, ) # Stamp start_time so save_all can record per-game duration. live.start_time = time.time() return live.session_id def _pick_partner(self, renderer: EnvRenderer, study: StudySession, idx: int) -> str: cfg = study.config is_warmup = idx < cfg.num_warmup pool = (cfg.warmup_partner_keys if is_warmup else cfg.real_partner_keys) if pool: return random.choice(pool) # Fallback: pick from full available_partners list (Johnny's _choose_agent equivalent). partners = [p.key for p in renderer.available_partners()] return random.choice(partners) if partners else renderer.default_partner_key() def purge_stale(self, ttl_seconds: int = 7200) -> int: cutoff = time.time() - ttl_seconds with self._lock: stale = [sid for sid, s in self._active.items() if s.last_seen_at < cutoff] for sid in stale: self._active.pop(sid) return len(stale) def _gen_id(prefix: str = "study_") -> str: import uuid return prefix + uuid.uuid4().hex[:12] def _gen_completion_code(length: int = 8) -> str: chars = string.ascii_uppercase + string.digits return "".join(random.choice(chars) for _ in range(length))