Spaces:
Runtime error
Runtime error
| """ | |
| SemanticTransitionDataset — PyTorch-compatible dataset for MBRL research. | |
| Loads trajectory JSONL files exported from CodeReviewEnv episodes and | |
| provides (state, action, reward, next_state, done) transitions for | |
| training Knowledge-Work World Models (KW-WM). | |
| Usage: | |
| from dataset import SemanticTransitionDataset | |
| ds = SemanticTransitionDataset("trajectories/") | |
| print(f"{len(ds)} transitions collected") | |
| # Each item is a dict with keys: | |
| # state_text: str — serialized observation (PR diff, context) | |
| # action_text: str — serialized action (type + params) | |
| # reward: float — grader reward for this transition | |
| # next_state_text: str — serialized next observation | |
| # done: bool — whether episode ended | |
| # task: str — easy|medium|hard | |
| # step: int — step number in episode | |
| # For embedding-based world models: | |
| from sentence_transformers import SentenceTransformer | |
| encoder = SentenceTransformer("all-MiniLM-L6-v2") | |
| s_enc = encoder.encode(ds[0]["state_text"]) | |
| # Train: MLP(s_enc, a_enc) → (s'_enc, r_pred) | |
| Research context: | |
| Standard MBRL benchmarks (Dreamer, MBPO, MuZero) assume vector state | |
| spaces with physics-based transitions. Text-based world models (Li et al., | |
| 2025) study synthetic text games; embodied SWMs (Berg et al., 2025) target | |
| robotics. CodeReviewEnv enables training the first **Knowledge-Work World | |
| Models (KW-WM)** — world models over structured professional text where | |
| T(s,a)→s' depends on professional judgment rather than physics or game rules. | |
| Open questions this dataset enables: | |
| 1. Does prediction error compound exponentially in knowledge-work spaces? | |
| 2. Does structured text provide natural error correction vs. continuous? | |
| 3. Can a KW-WM transfer across knowledge-work domains? | |
| """ | |
| import json | |
| import os | |
| from pathlib import Path | |
| from typing import Any, Dict, List, Optional | |
| class SemanticTransitionDataset: | |
| """ | |
| Loads trajectory JSONL files for training a Knowledge-Work World Model. | |
| Compatible with PyTorch Dataset interface (implements __len__ and __getitem__). | |
| Each trajectory file is a JSONL where each line is a transition dict. | |
| Args: | |
| trajectory_dir: Path to directory containing .jsonl trajectory files | |
| task_filter: Optional — only load trajectories for this task (easy|medium|hard) | |
| max_transitions: Optional — cap total transitions loaded (for memory) | |
| """ | |
| def __init__( | |
| self, | |
| trajectory_dir: str, | |
| task_filter: Optional[str] = None, | |
| max_transitions: Optional[int] = None, | |
| ): | |
| self.trajectory_dir = Path(trajectory_dir) | |
| self.transitions: List[Dict[str, Any]] = [] | |
| self._load(task_filter, max_transitions) | |
| def _load(self, task_filter: Optional[str], max_transitions: Optional[int]) -> None: | |
| """Load all .jsonl files from the trajectory directory.""" | |
| if not self.trajectory_dir.exists(): | |
| return | |
| for fpath in sorted(self.trajectory_dir.glob("*.jsonl")): | |
| with open(fpath) as f: | |
| for line in f: | |
| line = line.strip() | |
| if not line: | |
| continue | |
| try: | |
| transition = json.loads(line) | |
| except json.JSONDecodeError: | |
| continue | |
| # Apply task filter if specified | |
| if task_filter and transition.get("task") != task_filter: | |
| continue | |
| self.transitions.append(transition) | |
| if max_transitions and len(self.transitions) >= max_transitions: | |
| return | |
| def __len__(self) -> int: | |
| return len(self.transitions) | |
| def __getitem__(self, idx: int) -> Dict[str, Any]: | |
| """Return a single transition as a dict. | |
| Keys: | |
| state_text (str): Serialized observation text | |
| action_text (str): Serialized action string | |
| reward (float): Step reward | |
| next_state_text (str): Serialized next observation | |
| done (bool): Whether episode ended | |
| task (str): Task difficulty level | |
| step (int): Step number in episode | |
| """ | |
| t = self.transitions[idx] | |
| return { | |
| "state_text": t.get("state_text", json.dumps(t.get("state", {}))), | |
| "action_text": t.get("action_text", json.dumps(t.get("action", {}))), | |
| "reward": self._extract_reward(t), | |
| "next_state_text": t.get("next_state_text", json.dumps(t.get("next_state", {}))), | |
| "done": bool(t.get("done", False)), | |
| "task": t.get("task", "unknown"), | |
| "step": int(t.get("step", 0)), | |
| } | |
| def _extract_reward(t: Dict) -> float: | |
| """Extract reward as float, handling dict or float formats.""" | |
| r = t.get("reward", 0.0) | |
| if isinstance(r, dict): | |
| return float(r.get("value", 0.0)) | |
| try: | |
| return float(r) | |
| except (TypeError, ValueError): | |
| return 0.0 | |
| def get_episode(self, episode_id: str) -> List[Dict[str, Any]]: | |
| """Get all transitions from a specific episode.""" | |
| return [t for t in self.transitions if t.get("episode_id") == episode_id] | |
| def get_episodes(self) -> List[str]: | |
| """Get all unique episode IDs.""" | |
| return list(set(t.get("episode_id", "unknown") for t in self.transitions)) | |
| def stats(self) -> Dict[str, Any]: | |
| """Summary statistics for the loaded dataset.""" | |
| episodes = self.get_episodes() | |
| rewards = [self._extract_reward(t) for t in self.transitions] | |
| tasks = {} | |
| for t in self.transitions: | |
| task = t.get("task", "unknown") | |
| tasks[task] = tasks.get(task, 0) + 1 | |
| return { | |
| "total_transitions": len(self.transitions), | |
| "total_episodes": len(episodes), | |
| "task_distribution": tasks, | |
| "reward_mean": sum(rewards) / len(rewards) if rewards else 0.0, | |
| "reward_min": min(rewards) if rewards else 0.0, | |
| "reward_max": max(rewards) if rewards else 0.0, | |
| } | |
| def to_pytorch(self): | |
| """Convert to a PyTorch-compatible dataset (requires torch).""" | |
| try: | |
| import torch | |
| from torch.utils.data import Dataset as TorchDataset | |
| parent = self | |
| class _TorchWrapper(TorchDataset): | |
| def __len__(self): | |
| return len(parent) | |
| def __getitem__(self, idx): | |
| item = parent[idx] | |
| return { | |
| "state_text": item["state_text"], | |
| "action_text": item["action_text"], | |
| "reward": torch.tensor(item["reward"], dtype=torch.float32), | |
| "next_state_text": item["next_state_text"], | |
| "done": torch.tensor(item["done"], dtype=torch.bool), | |
| } | |
| return _TorchWrapper() | |
| except ImportError: | |
| raise ImportError("PyTorch is required for to_pytorch(). Install with: pip install torch") | |