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

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

@@ -1,431 +0,0 @@
-"""Visualization tools for NetHack trajectories."""
-
-import json
-from datetime import datetime
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
-
-import matplotlib.animation as animation
-import matplotlib.pyplot as plt
-import numpy as np
-import seaborn as sns
-from matplotlib.patches import Rectangle
-
-try:
-    from PIL import Image, ImageDraw, ImageFont
-
-    HAS_PIL = True
-except ImportError:
-    HAS_PIL = False
-    print("Warning: PIL not available. Some visualization features will be limited.")
-
-
-class NetHackVisualizer:
-    """Visualize NetHack game states and trajectories."""
-
-    # Character to color mapping for ASCII visualization
-    CHAR_COLORS = {
-        "@": "#FF0000",  # Player - red
-        "d": "#8B4513",  # Dog/pet - brown
-        "f": "#FFA500",  # Cat - orange
-        ">": "#0000FF",  # Stairs down - blue
-        "<": "#00FF00",  # Stairs up - green
-        ".": "#D3D3D3",  # Floor - light gray
-        "#": "#696969",  # Wall - dark gray
-        "+": "#8B4513",  # Door - brown
-        "-": "#8B4513",  # Door - brown
-        "|": "#8B4513",  # Door - brown
-        "{": "#00CED1",  # Fountain - dark turquoise
-        "}": "#4682B4",  # Pool - steel blue
-        "^": "#FF1493",  # Trap - deep pink
-        "%": "#FFD700",  # Food - gold
-        "!": "#FF69B4",  # Potion - hot pink
-        "?": "#DDA0DD",  # Scroll - plum
-        "/": "#9370DB",  # Wand - medium purple
-        "=": "#FFD700",  # Ring - gold
-        '"': "#FF4500",  # Amulet - orange red
-        "[": "#C0C0C0",  # Armor - silver
-        ")": "#A9A9A9",  # Weapon - dark gray
-        "*": "#FFFF00",  # Gold/gem - yellow
-        "$": "#FFD700",  # Gold - gold
-        "`": "#8B4513",  # Boulder/statue - brown
-    }
-
-    # Default color for unknown characters
-    DEFAULT_COLOR = "#FFFFFF"
-
-    def __init__(self, cell_size: int = 10, font_size: int = 8):
-        self.cell_size = cell_size
-        self.font_size = font_size
-
-    def ascii_to_image(
-        self, ascii_map: str, highlight_pos: Optional[Tuple[int, int]] = None
-    ) -> np.ndarray:
-        """Convert ASCII map to colored image."""
-        if not HAS_PIL:
-            return self._simple_ascii_to_image(ascii_map, highlight_pos)
-
-        lines = ascii_map.strip().split("\n")
-        height = len(lines)
-        width = max(len(line) for line in lines) if lines else 0
-
-        # Create image
-        img_width = width * self.cell_size
-        img_height = height * self.cell_size
-        image = Image.new("RGB", (img_width, img_height), color="black")
-        draw = ImageDraw.Draw(image)
-
-        # Try to load a monospace font
-        try:
-            font = ImageFont.truetype("/System/Library/Fonts/Monaco.ttf", self.font_size)
-        except:
-            font = ImageFont.load_default()
-
-        # Draw each character
-        for y, line in enumerate(lines):
-            for x, char in enumerate(line):
-                if char == " ":
-                    continue
-
-                # Get color for character
-                color = self.CHAR_COLORS.get(char, self.DEFAULT_COLOR)
-
-                # Highlight player position
-                if highlight_pos and (x, y) == highlight_pos:
-                    # Draw background
-                    draw.rectangle(
-                        [
-                            x * self.cell_size,
-                            y * self.cell_size,
-                            (x + 1) * self.cell_size,
-                            (y + 1) * self.cell_size,
-                        ],
-                        fill="yellow",
-                    )
-
-                # Draw character
-                draw.text(
-                    (x * self.cell_size + 2, y * self.cell_size),
-                    char,
-                    fill=color,
-                    font=font,
-                )
-
-        return np.array(image)
-
-    def _simple_ascii_to_image(
-        self, ascii_map: str, highlight_pos: Optional[Tuple[int, int]] = None
-    ) -> np.ndarray:
-        """Simple ASCII to image conversion without PIL."""
-        lines = ascii_map.strip().split("\n")
-        height = len(lines)
-        width = max(len(line) for line in lines) if lines else 0
-
-        # Create RGB image
-        image = np.zeros((height * self.cell_size, width * self.cell_size, 3), dtype=np.uint8)
-
-        # Simple character to grayscale mapping
-        char_values = {
-            "@": 255,  # Player - white
-            "#": 64,  # Wall - dark gray
-            ".": 128,  # Floor - gray
-            ">": 200,  # Stairs - light
-            "<": 200,  # Stairs - light
-            " ": 0,  # Empty - black
-        }
-
-        for y, line in enumerate(lines):
-            for x, char in enumerate(line):
-                value = char_values.get(char, 100)
-                y_start = y * self.cell_size
-                y_end = (y + 1) * self.cell_size
-                x_start = x * self.cell_size
-                x_end = (x + 1) * self.cell_size
-
-                if highlight_pos and (x, y) == highlight_pos:
-                    # Highlight player in red
-                    image[y_start:y_end, x_start:x_end] = [255, 0, 0]
-                else:
-                    image[y_start:y_end, x_start:x_end] = [value, value, value]
-
-        return image
-
-    def create_frame_image(
-        self, observation: Dict[str, Any], include_stats: bool = True
-    ) -> np.ndarray:
-        """Create a single frame image from observation."""
-        # Standard NetHack terminal size
-        TERM_WIDTH = 80
-        TERM_HEIGHT = 24
-
-        if HAS_PIL:
-            return self._create_terminal_view(observation, TERM_WIDTH, TERM_HEIGHT)
-        else:
-            # Fallback to simple view
-            ascii_map = observation.get("ascii_map", "")
-            player_stats = observation.get("player_stats", {})
-            player_pos = (player_stats.get("x", 0), player_stats.get("y", 0))
-            return self.ascii_to_image(ascii_map, player_pos)
-
-    def _create_terminal_view(
-        self, observation: Dict[str, Any], width: int = 80, height: int = 24
-    ) -> np.ndarray:
-        """Create a full terminal-style NetHack view."""
-        # Create black background
-        img_width = width * self.cell_size
-        img_height = height * self.cell_size
-        image = Image.new("RGB", (img_width, img_height), color="black")
-        draw = ImageDraw.Draw(image)
-
-        # Try to load a monospace font
-        try:
-            font = ImageFont.truetype("/System/Library/Fonts/Monaco.ttf", self.font_size)
-        except:
-            try:
-                font = ImageFont.truetype(
-                    "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf",
-                    self.font_size,
-                )
-            except:
-                font = ImageFont.load_default()
-
-        # Get game data
-        ascii_map = observation.get("ascii_map", "")
-        player_stats = observation.get("player_stats", {})
-        message = observation.get("message", "").strip()
-
-        # Draw the map (first 21 lines)
-        lines = ascii_map.split("\n")
-        for y, line in enumerate(lines[:21]):
-            for x, char in enumerate(line[:width]):
-                if char == " ":
-                    continue
-
-                # Get color for character
-                color = self.CHAR_COLORS.get(char, self.DEFAULT_COLOR)
-
-                # Highlight player position
-                if (x, y) == (player_stats.get("x", -1), player_stats.get("y", -1)):
-                    # Draw yellow background for player
-                    draw.rectangle(
-                        [
-                            x * self.cell_size,
-                            y * self.cell_size,
-                            (x + 1) * self.cell_size,
-                            (y + 1) * self.cell_size,
-                        ],
-                        fill="#333300",
-                    )
-
-                # Draw character
-                draw.text(
-                    (x * self.cell_size + 1, y * self.cell_size),
-                    char,
-                    fill=color,
-                    font=font,
-                )
-
-        # Draw separator line (line 21)
-        draw.line(
-            [(0, 21 * self.cell_size), (img_width, 21 * self.cell_size)],
-            fill="#666666",
-            width=1,
-        )
-
-        # Draw status line (line 22)
-        status_y = 22 * self.cell_size
-        character_name = observation.get("character_name", "Agent")
-        character_role = observation.get("character_role", "Adventurer")
-
-        # Format status line like real NetHack
-        status_parts = [
-            f"{character_name} the {character_role}",
-            f"St:{player_stats.get('strength', 10)}",
-            f"Dx:{player_stats.get('dexterity', 10)}",
-            f"Co:{player_stats.get('constitution', 10)}",
-            f"In:{player_stats.get('intelligence', 10)}",
-            f"Wi:{player_stats.get('wisdom', 10)}",
-            f"Ch:{player_stats.get('charisma', 10)}",
-            "Neutral",  # Alignment
-        ]
-
-        status_line = "  ".join(status_parts)
-        draw.text((5, status_y), status_line, fill="white", font=font)
-
-        # Draw second status line (line 23)
-        status2_y = 23 * self.cell_size
-        dlvl = player_stats.get("depth", 1)
-        gold = player_stats.get("gold", 0)
-        hp = player_stats.get("hp", 10)
-        max_hp = player_stats.get("max_hp", 10)
-        pw = player_stats.get("energy", 0)
-        max_pw = player_stats.get("max_energy", 0)
-        ac = player_stats.get("ac", 10)
-        xp = player_stats.get("experience_level", 1)
-
-        status2_parts = [
-            f"Dlvl:{dlvl}",
-            f"$:{gold}",
-            f"HP:{hp}({max_hp})",
-            f"Pw:{pw}({max_pw})",
-            f"AC:{ac}",
-            f"Xp:{xp}",
-        ]
-
-        # Add turn count if available
-        if "turn_count" in observation:
-            status2_parts.append(f"T:{observation['turn_count']}")
-
-        status2_line = "  ".join(status2_parts)
-        draw.text((5, status2_y), status2_line, fill="white", font=font)
-
-        # Draw message at top if present
-        if message:
-            # Clear message area (line 0)
-            draw.rectangle([0, 0, img_width, self.cell_size], fill="black")
-            draw.text((5, 0), message[: width - 1], fill="white", font=font)
-
-        return np.array(image)
-
-    def create_trajectory_video(
-        self,
-        frames: List[Dict[str, Any]],
-        output_path: str,
-        fps: int = 4,
-        include_stats: bool = True,
-    ) -> str:
-        """Create a video from trajectory frames."""
-        if not frames:
-            raise ValueError("No frames to create video from")
-
-        # Create figure and axis
-        fig, ax = plt.subplots(figsize=(10, 12))
-        ax.axis("off")
-
-        # Create first frame
-        first_img = self.create_frame_image(frames[0]["observation"], include_stats)
-        im = ax.imshow(first_img)
-
-        def animate(i):
-            if i < len(frames):
-                img = self.create_frame_image(frames[i]["observation"], include_stats)
-                im.set_array(img)
-                ax.set_title(f"Step {i}: {frames[i]['action']}")
-            return [im]
-
-        # Create animation
-        anim = animation.FuncAnimation(
-            fig, animate, frames=len(frames), interval=1000 / fps, blit=True
-        )
-
-        # Save as video
-        output_path = Path(output_path)
-        output_path.parent.mkdir(parents=True, exist_ok=True)
-
-        if output_path.suffix == ".gif":
-            anim.save(str(output_path), writer="pillow", fps=fps)
-        else:
-            anim.save(str(output_path), writer="ffmpeg", fps=fps)
-
-        plt.close()
-
-        return str(output_path)
-
-    def plot_trajectory_stats(
-        self, frames: List[Dict[str, Any]], output_path: Optional[str] = None
-    ):
-        """Plot statistics from a trajectory."""
-        if not frames:
-            return
-
-        # Extract data
-        steps = []
-        rewards = []
-        depths = []
-        hps = []
-        positions_x = []
-        positions_y = []
-
-        cumulative_reward = 0
-        for i, frame in enumerate(frames):
-            steps.append(i)
-            cumulative_reward += frame["reward"]
-            rewards.append(cumulative_reward)
-
-            stats = frame["observation"].get("player_stats", {})
-            depths.append(stats.get("depth", 1))
-            hps.append(stats.get("hp", 0))
-            positions_x.append(stats.get("x", 0))
-            positions_y.append(stats.get("y", 0))
-
-        # Create subplots
-        fig, axes = plt.subplots(2, 2, figsize=(12, 10))
-
-        # Cumulative reward
-        axes[0, 0].plot(steps, rewards, "b-")
-        axes[0, 0].set_xlabel("Step")
-        axes[0, 0].set_ylabel("Cumulative Reward")
-        axes[0, 0].set_title("Reward Progress")
-        axes[0, 0].grid(True)
-
-        # Depth progression
-        axes[0, 1].plot(steps, depths, "g-")
-        axes[0, 1].set_xlabel("Step")
-        axes[0, 1].set_ylabel("Dungeon Depth")
-        axes[0, 1].set_title("Depth Exploration")
-        axes[0, 1].grid(True)
-
-        # HP over time
-        axes[1, 0].plot(steps, hps, "r-")
-        axes[1, 0].set_xlabel("Step")
-        axes[1, 0].set_ylabel("Hit Points")
-        axes[1, 0].set_title("Health Over Time")
-        axes[1, 0].grid(True)
-
-        # Position heatmap
-        axes[1, 1].scatter(positions_x, positions_y, c=steps, cmap="viridis", s=1)
-        axes[1, 1].set_xlabel("X Position")
-        axes[1, 1].set_ylabel("Y Position")
-        axes[1, 1].set_title("Movement Pattern")
-        axes[1, 1].invert_yaxis()  # Invert Y axis for proper orientation
-
-        plt.tight_layout()
-
-        if output_path:
-            plt.savefig(output_path, dpi=150)
-            plt.close()
-        else:
-            plt.show()
-
-    def plot_action_distribution(
-        self, frames: List[Dict[str, Any]], output_path: Optional[str] = None
-    ):
-        """Plot distribution of actions taken."""
-        if not frames:
-            return
-
-        # Count actions
-        action_counts = {}
-        for frame in frames:
-            action = frame["action"]
-            action_counts[action] = action_counts.get(action, 0) + 1
-
-        # Sort by frequency
-        actions = sorted(action_counts.items(), key=lambda x: x[1], reverse=True)
-
-        # Create bar plot
-        plt.figure(figsize=(12, 6))
-        actions_list, counts = zip(*actions[:20])  # Top 20 actions
-
-        plt.bar(range(len(actions_list)), counts)
-        plt.xticks(range(len(actions_list)), actions_list, rotation=45, ha="right")
-        plt.xlabel("Action")
-        plt.ylabel("Count")
-        plt.title("Action Distribution (Top 20)")
-        plt.tight_layout()
-
-        if output_path:
-            plt.savefig(output_path, dpi=150)
-            plt.close()
-        else:
-            plt.show()