synth-ai 0.2.14__py3-none-any.whl → 0.4.4__py3-none-any.whl

synth_ai/environments/examples/nethack/helpers/trajectory_recorder.py

@@ -1,269 +0,0 @@
-"""Trajectory recording and replay functionality for NetHack."""
-
-import base64
-import gzip
-import json
-import pickle
-from dataclasses import asdict, dataclass
-from datetime import datetime
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
-
-import numpy as np
-
-
-@dataclass
-class TrajectoryFrame:
-    """Single frame in a trajectory."""
-
-    step: int
-    action: str
-    observation: Dict[str, Any]
-    reward: float
-    done: bool
-    info: Dict[str, Any]
-    timestamp: float
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert to serializable dict."""
-        d = asdict(self)
-        # Handle numpy arrays in observation
-        if "observation" in d and d["observation"]:
-            d["observation"] = self._serialize_observation(d["observation"])
-        return d
-
-    def _serialize_observation(self, obs: Dict[str, Any]) -> Dict[str, Any]:
-        """Serialize observation, converting numpy arrays to lists."""
-        serialized = {}
-        for key, value in obs.items():
-            if isinstance(value, np.ndarray):
-                serialized[key] = {
-                    "type": "ndarray",
-                    "data": value.tolist(),
-                    "dtype": str(value.dtype),
-                    "shape": value.shape,
-                }
-            elif isinstance(value, dict):
-                serialized[key] = self._serialize_observation(value)
-            elif isinstance(value, (list, tuple)) and len(value) > 0 and isinstance(value[0], dict):
-                serialized[key] = [
-                    self._serialize_observation(item) if isinstance(item, dict) else item
-                    for item in value
-                ]
-            else:
-                serialized[key] = value
-        return serialized
-
-    @classmethod
-    def from_dict(cls, d: Dict[str, Any]) -> "TrajectoryFrame":
-        """Reconstruct from dict."""
-        if "observation" in d and d["observation"]:
-            d["observation"] = cls._deserialize_observation(d["observation"])
-        return cls(**d)
-
-    @staticmethod
-    def _deserialize_observation(obs: Dict[str, Any]) -> Dict[str, Any]:
-        """Deserialize observation, converting lists back to numpy arrays."""
-        deserialized = {}
-        for key, value in obs.items():
-            if isinstance(value, dict) and value.get("type") == "ndarray":
-                deserialized[key] = np.array(value["data"], dtype=value["dtype"]).reshape(
-                    value["shape"]
-                )
-            elif isinstance(value, dict):
-                deserialized[key] = TrajectoryFrame._deserialize_observation(value)
-            elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], dict):
-                deserialized[key] = [
-                    TrajectoryFrame._deserialize_observation(item)
-                    if isinstance(item, dict)
-                    else item
-                    for item in value
-                ]
-            else:
-                deserialized[key] = value
-        return deserialized
-
-
-@dataclass
-class TrajectoryMetadata:
-    """Metadata for a trajectory."""
-
-    trajectory_id: str
-    character_role: str
-    task_id: Optional[str]
-    start_time: datetime
-    end_time: Optional[datetime]
-    total_steps: int
-    total_reward: float
-    final_status: str  # 'completed', 'died', 'quit', 'truncated'
-    max_depth_reached: int
-    achievements: Dict[str, bool]
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert to serializable dict."""
-        d = asdict(self)
-        d["start_time"] = self.start_time.isoformat()
-        if self.end_time:
-            d["end_time"] = self.end_time.isoformat()
-        return d
-
-    @classmethod
-    def from_dict(cls, d: Dict[str, Any]) -> "TrajectoryMetadata":
-        """Reconstruct from dict."""
-        d["start_time"] = datetime.fromisoformat(d["start_time"])
-        if d.get("end_time"):
-            d["end_time"] = datetime.fromisoformat(d["end_time"])
-        return cls(**d)
-
-
-class TrajectoryRecorder:
-    """Records and saves NetHack game trajectories."""
-
-    def __init__(self, save_dir: str = "temp/nethack_trajectories"):
-        self.save_dir = Path(save_dir)
-        self.save_dir.mkdir(parents=True, exist_ok=True)
-
-        self.frames: List[TrajectoryFrame] = []
-        self.metadata: Optional[TrajectoryMetadata] = None
-        self.is_recording = False
-        self.current_step = 0
-        self.total_reward = 0.0
-        self.max_depth = 1
-
-    def start_recording(self, character_role: str, task_id: Optional[str] = None) -> str:
-        """Start recording a new trajectory."""
-        trajectory_id = datetime.now().strftime("%Y%m%d_%H%M%S_") + character_role
-
-        self.metadata = TrajectoryMetadata(
-            trajectory_id=trajectory_id,
-            character_role=character_role,
-            task_id=task_id,
-            start_time=datetime.now(),
-            end_time=None,
-            total_steps=0,
-            total_reward=0.0,
-            final_status="in_progress",
-            max_depth_reached=1,
-            achievements={},
-        )
-
-        self.frames = []
-        self.is_recording = True
-        self.current_step = 0
-        self.total_reward = 0.0
-        self.max_depth = 1
-
-        return trajectory_id
-
-    def record_step(
-        self,
-        action: str,
-        observation: Dict[str, Any],
-        reward: float,
-        done: bool,
-        info: Dict[str, Any],
-    ):
-        """Record a single step."""
-        if not self.is_recording:
-            raise ValueError("Recording not started. Call start_recording first.")
-
-        frame = TrajectoryFrame(
-            step=self.current_step,
-            action=action,
-            observation=observation.copy(),
-            reward=reward,
-            done=done,
-            info=info.copy(),
-            timestamp=datetime.now().timestamp(),
-        )
-
-        self.frames.append(frame)
-        self.current_step += 1
-        self.total_reward += reward
-
-        # Update max depth if available
-        if "player_stats" in observation:
-            depth = observation["player_stats"].get("depth", 1)
-            self.max_depth = max(self.max_depth, depth)
-
-    def stop_recording(
-        self,
-        final_status: str = "completed",
-        achievements: Optional[Dict[str, bool]] = None,
-    ):
-        """Stop recording and finalize metadata."""
-        if not self.is_recording:
-            return
-
-        self.is_recording = False
-
-        if self.metadata:
-            self.metadata.end_time = datetime.now()
-            self.metadata.total_steps = self.current_step
-            self.metadata.total_reward = self.total_reward
-            self.metadata.final_status = final_status
-            self.metadata.max_depth_reached = self.max_depth
-            if achievements:
-                self.metadata.achievements = achievements
-
-    def save_trajectory(self, filename: Optional[str] = None) -> str:
-        """Save trajectory to disk."""
-        if not self.metadata:
-            raise ValueError("No trajectory to save")
-
-        if filename is None:
-            filename = f"{self.metadata.trajectory_id}.trajectory.gz"
-
-        filepath = self.save_dir / filename
-
-        trajectory_data = {
-            "metadata": self.metadata.to_dict(),
-            "frames": [frame.to_dict() for frame in self.frames],
-        }
-
-        # Save as compressed JSON
-        with gzip.open(filepath, "wt", encoding="utf-8") as f:
-            json.dump(trajectory_data, f, indent=2)
-
-        # Also save a quick info file
-        info_file = self.save_dir / f"{self.metadata.trajectory_id}.info.json"
-        with open(info_file, "w") as f:
-            json.dump(self.metadata.to_dict(), f, indent=2)
-
-        return str(filepath)
-
-    @classmethod
-    def load_trajectory(
-        cls, filepath: str
-    ) -> Tuple["TrajectoryRecorder", TrajectoryMetadata, List[TrajectoryFrame]]:
-        """Load a trajectory from disk."""
-        recorder = cls()
-
-        with gzip.open(filepath, "rt", encoding="utf-8") as f:
-            data = json.load(f)
-
-        metadata = TrajectoryMetadata.from_dict(data["metadata"])
-        frames = [TrajectoryFrame.from_dict(frame) for frame in data["frames"]]
-
-        recorder.metadata = metadata
-        recorder.frames = frames
-
-        return recorder, metadata, frames
-
-    def get_summary(self) -> Dict[str, Any]:
-        """Get summary statistics of the current trajectory."""
-        if not self.frames:
-            return {}
-
-        actions_taken = {}
-        for frame in self.frames:
-            actions_taken[frame.action] = actions_taken.get(frame.action, 0) + 1
-
-        return {
-            "total_steps": len(self.frames),
-            "total_reward": self.total_reward,
-            "max_depth": self.max_depth,
-            "actions_distribution": actions_taken,
-            "unique_actions": len(actions_taken),
-            "average_reward_per_step": self.total_reward / len(self.frames) if self.frames else 0,
-        }

synth_ai/environments/examples/nethack/helpers/visualization/replay_viewer.py

@@ -1,308 +0,0 @@
-"""Interactive replay viewer for NetHack trajectories."""
-
-import gzip
-import json
-import os
-import sys
-from datetime import datetime
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
-
-# Add parent directory to path for imports
-sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent.parent))
-
-from src.synth_env.examples.nethack.helpers.trajectory_recorder import (
-    TrajectoryFrame,
-    TrajectoryRecorder,
-)
-from src.synth_env.examples.nethack.helpers.visualization.visualizer import (
-    NetHackVisualizer,
-)
-
-
-class ReplayViewer:
-    """Interactive viewer for NetHack trajectory replays."""
-
-    def __init__(self, trajectory_path: str):
-        """Initialize replay viewer with a trajectory file."""
-        self.trajectory_path = Path(trajectory_path)
-
-        # Load trajectory
-        self.recorder, self.metadata, self.frames = TrajectoryRecorder.load_trajectory(
-            str(trajectory_path)
-        )
-
-        # Initialize visualizer
-        self.visualizer = NetHackVisualizer()
-
-        # Playback state
-        self.current_frame = 0
-        self.is_playing = False
-        self.playback_speed = 1.0  # Frames per second
-
-        print(f"Loaded trajectory: {self.metadata.trajectory_id}")
-        print(f"Character: {self.metadata.character_role}")
-        print(f"Total steps: {self.metadata.total_steps}")
-        print(f"Total reward: {self.metadata.total_reward:.2f}")
-        print(f"Final status: {self.metadata.final_status}")
-        print(f"Max depth reached: {self.metadata.max_depth_reached}")
-
-    def show_frame(self, frame_idx: int):
-        """Display a specific frame."""
-        if 0 <= frame_idx < len(self.frames):
-            frame = self.frames[frame_idx]
-            obs = frame.observation
-
-            print(f"\n{'=' * 80}")
-            print(f"Frame {frame_idx}/{len(self.frames) - 1} | Step: {frame.step}")
-            print(f"Action: {frame.action}")
-            print(f"Reward: {frame.reward:+.2f}")
-
-            # Show message
-            message = obs.get("message", "").strip()
-            if message:
-                print(f"Message: {message}")
-
-            # Show stats
-            stats = obs.get("player_stats", {})
-            print(
-                f"Position: ({stats.get('x', 0)}, {stats.get('y', 0)}) | "
-                f"HP: {stats.get('hp', 0)}/{stats.get('max_hp', 0)} | "
-                f"Level: {stats.get('experience_level', 1)} | "
-                f"Depth: {stats.get('depth', 1)} | "
-                f"Gold: {stats.get('gold', 0)}"
-            )
-
-            # Show map
-            ascii_map = obs.get("ascii_map", "")
-            if ascii_map:
-                lines = ascii_map.split("\n")
-                px, py = stats.get("x", 0), stats.get("y", 0)
-
-                # Show area around player
-                print("\nMap view:")
-                for y in range(max(0, py - 10), min(len(lines), py + 11)):
-                    if 0 <= y < len(lines):
-                        line = lines[y]
-                        start = max(0, px - 20)
-                        end = min(len(line), px + 21)
-                        if y == py:
-                            print(f">>> {line[start:end]} <<<")
-                        else:
-                            print(f"    {line[start:end]}")
-
-    def interactive_replay(self):
-        """Run interactive replay session."""
-        print("\n=== Interactive Replay Viewer ===")
-        print("Commands:")
-        print("  n/next - Next frame")
-        print("  p/prev - Previous frame")
-        print("  g <num> - Go to frame number")
-        print("  f/first - Go to first frame")
-        print("  l/last - Go to last frame")
-        print("  i/info - Show trajectory info")
-        print("  s/search <action> - Find frames with action")
-        print("  export <type> - Export (video/stats/actions)")
-        print("  q/quit - Exit viewer")
-
-        # Show first frame
-        self.show_frame(0)
-
-        while True:
-            try:
-                cmd = input(f"\nFrame {self.current_frame}> ").strip().lower()
-
-                if cmd in ["q", "quit"]:
-                    break
-
-                elif cmd in ["n", "next"]:
-                    if self.current_frame < len(self.frames) - 1:
-                        self.current_frame += 1
-                        self.show_frame(self.current_frame)
-                    else:
-                        print("Already at last frame")
-
-                elif cmd in ["p", "prev"]:
-                    if self.current_frame > 0:
-                        self.current_frame -= 1
-                        self.show_frame(self.current_frame)
-                    else:
-                        print("Already at first frame")
-
-                elif cmd.startswith("g "):
-                    try:
-                        frame_num = int(cmd.split()[1])
-                        if 0 <= frame_num < len(self.frames):
-                            self.current_frame = frame_num
-                            self.show_frame(self.current_frame)
-                        else:
-                            print(f"Frame number must be between 0 and {len(self.frames) - 1}")
-                    except (ValueError, IndexError):
-                        print("Usage: g <frame_number>")
-
-                elif cmd in ["f", "first"]:
-                    self.current_frame = 0
-                    self.show_frame(self.current_frame)
-
-                elif cmd in ["l", "last"]:
-                    self.current_frame = len(self.frames) - 1
-                    self.show_frame(self.current_frame)
-
-                elif cmd in ["i", "info"]:
-                    self.show_trajectory_info()
-
-                elif cmd.startswith("s ") or cmd.startswith("search "):
-                    action = " ".join(cmd.split()[1:])
-                    self.search_action(action)
-
-                elif cmd.startswith("export"):
-                    parts = cmd.split()
-                    if len(parts) > 1:
-                        self.export_trajectory(parts[1])
-                    else:
-                        print("Usage: export <video|stats|actions>")
-
-                else:
-                    print("Unknown command. Type 'q' to quit.")
-
-            except KeyboardInterrupt:
-                print("\nUse 'q' to quit")
-            except Exception as e:
-                print(f"Error: {e}")
-
-    def show_trajectory_info(self):
-        """Display trajectory information and statistics."""
-        print(f"\n=== Trajectory Information ===")
-        print(f"ID: {self.metadata.trajectory_id}")
-        print(f"Character: {self.metadata.character_role}")
-        print(f"Task ID: {self.metadata.task_id or 'N/A'}")
-        print(f"Start time: {self.metadata.start_time}")
-        print(f"End time: {self.metadata.end_time}")
-        print(
-            f"Duration: {(self.metadata.end_time - self.metadata.start_time).total_seconds():.1f} seconds"
-        )
-        print(f"Total steps: {self.metadata.total_steps}")
-        print(f"Total reward: {self.metadata.total_reward:.2f}")
-        print(f"Average reward/step: {self.metadata.total_reward / self.metadata.total_steps:.4f}")
-        print(f"Final status: {self.metadata.final_status}")
-        print(f"Max depth: {self.metadata.max_depth_reached}")
-
-        # Action distribution
-        action_counts = {}
-        for frame in self.frames:
-            action_counts[frame.action] = action_counts.get(frame.action, 0) + 1
-
-        print(f"\nTop 10 actions:")
-        for action, count in sorted(action_counts.items(), key=lambda x: x[1], reverse=True)[:10]:
-            print(f"  {action}: {count} ({count / len(self.frames) * 100:.1f}%)")
-
-    def search_action(self, action: str):
-        """Search for frames containing specific action."""
-        matches = []
-        for i, frame in enumerate(self.frames):
-            if action.lower() in frame.action.lower():
-                matches.append(i)
-
-        if matches:
-            print(f"Found {len(matches)} frames with action '{action}':")
-            for i, frame_idx in enumerate(matches[:10]):  # Show first 10
-                frame = self.frames[frame_idx]
-                print(f"  Frame {frame_idx}: {frame.action} (reward: {frame.reward:+.2f})")
-
-            if len(matches) > 10:
-                print(f"  ... and {len(matches) - 10} more")
-        else:
-            print(f"No frames found with action '{action}'")
-
-    def export_trajectory(self, export_type: str):
-        """Export trajectory in various formats."""
-        output_dir = self.trajectory_path.parent / "exports"
-        output_dir.mkdir(exist_ok=True)
-
-        if export_type == "video":
-            print("Creating video...")
-            output_path = output_dir / f"{self.metadata.trajectory_id}.mp4"
-
-            # Convert frames to format expected by visualizer
-            vis_frames = []
-            for frame in self.frames:
-                vis_frames.append({"action": frame.action, "observation": frame.observation})
-
-            try:
-                video_path = self.visualizer.create_trajectory_video(
-                    vis_frames, str(output_path), fps=4, include_stats=True
-                )
-                print(f"Video saved to: {video_path}")
-            except Exception as e:
-                print(f"Error creating video: {e}")
-                print("Make sure ffmpeg is installed for video export")
-
-        elif export_type == "stats":
-            print("Creating statistics plots...")
-            stats_path = output_dir / f"{self.metadata.trajectory_id}_stats.png"
-            action_path = output_dir / f"{self.metadata.trajectory_id}_actions.png"
-
-            # Convert frames for visualizer
-            vis_frames = []
-            for frame in self.frames:
-                vis_frames.append(
-                    {
-                        "action": frame.action,
-                        "observation": frame.observation,
-                        "reward": frame.reward,
-                    }
-                )
-
-            self.visualizer.plot_trajectory_stats(vis_frames, str(stats_path))
-            self.visualizer.plot_action_distribution(vis_frames, str(action_path))
-            print(f"Stats saved to: {stats_path}")
-            print(f"Action distribution saved to: {action_path}")
-
-        elif export_type == "actions":
-            print("Exporting action sequence...")
-            actions_path = output_dir / f"{self.metadata.trajectory_id}_actions.txt"
-
-            with open(actions_path, "w") as f:
-                f.write(f"# Trajectory: {self.metadata.trajectory_id}\n")
-                f.write(f"# Character: {self.metadata.character_role}\n")
-                f.write(f"# Total steps: {self.metadata.total_steps}\n")
-                f.write(f"# Total reward: {self.metadata.total_reward}\n\n")
-
-                for frame in self.frames:
-                    f.write(f"{frame.step:4d}: {frame.action:15s} (reward: {frame.reward:+.2f})\n")
-
-            print(f"Action sequence saved to: {actions_path}")
-
-        else:
-            print(f"Unknown export type: {export_type}")
-            print("Available types: video, stats, actions")
-
-
-def main():
-    """Main entry point for replay viewer."""
-    import argparse
-
-    parser = argparse.ArgumentParser(description="NetHack Trajectory Replay Viewer")
-    parser.add_argument("trajectory", help="Path to trajectory file (.trajectory.gz)")
-    parser.add_argument(
-        "--export",
-        choices=["video", "stats", "actions"],
-        help="Export trajectory without interactive mode",
-    )
-
-    args = parser.parse_args()
-
-    if not Path(args.trajectory).exists():
-        print(f"Error: Trajectory file not found: {args.trajectory}")
-        sys.exit(1)
-
-    viewer = ReplayViewer(args.trajectory)
-
-    if args.export:
-        viewer.export_trajectory(args.export)
-    else:
-        viewer.interactive_replay()
-
-
-if __name__ == "__main__":
-    main()