evaluation.py

import argparse
import logging
import os
import sys
import time
from datetime import datetime
from typing import Dict, Optional

import gymnasium as gym
import metaworld
import numpy as np
from agent import RLAgent

from torch.utils.tensorboard import SummaryWriter


class AgentEvaluator:
    """
    Evaluator for running and assessing the agent in MetaWorld environments.
    Includes TensorBoard logging for performance monitoring.
    """

    def __init__(
        self,
        task_name: str = "reach-v3",
        render_mode: str = "human",
        max_episodes: int = 5,
        max_steps_per_episode: int = 200,
        seed: Optional[int] = None,
        use_tensorboard: bool = True,
        log_dir: Optional[str] = None,
    ):
        """
        Initialize the evaluator.

        Args:
            task_name: Name of the MetaWorld task to run
            render_mode: Rendering mode ("human" for GUI, "rgb_array" for headless)
            max_episodes: Maximum number of episodes to run
            max_steps_per_episode: Maximum steps per episode
            seed: Random seed for reproducibility
            use_tensorboard: Whether to enable TensorBoard logging
            log_dir: Directory for TensorBoard logs (auto-generated if None)
        """
        self.task_name = task_name
        self.render_mode = render_mode
        self.max_episodes = max_episodes
        self.max_steps_per_episode = max_steps_per_episode
        self.seed = seed or np.random.randint(0, 1000000)
        self.use_tensorboard = use_tensorboard

        self.logger = logging.getLogger(__name__)
        self.env = None
        self.agent = None

        # Statistics tracking
        self.episode_rewards = []
        self.episode_lengths = []
        self.success_rate = 0.0

        # TensorBoard setup
        self.tb_writer = None
        if self.use_tensorboard:
            if log_dir is None:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                log_dir = f"runs/{self.task_name}_{timestamp}"

            os.makedirs(log_dir, exist_ok=True)
            self.tb_writer = SummaryWriter(log_dir)
            self.logger.info(f"TensorBoard logging enabled: {log_dir}")
            self.logger.info(f"View logs with: tensorboard --logdir {log_dir}")

        """
        Initialize the evaluator.

        Args:
            task_name: Name of the MetaWorld task to run
            render_mode: Rendering mode ("human" for GUI, "rgb_array" for headless)
            max_episodes: Maximum number of episodes to run
            max_steps_per_episode: Maximum steps per episode
            seed: Random seed for reproducibility
            use_tensorboard: Whether to enable TensorBoard logging
            log_dir: Directory for TensorBoard logs (auto-generated if None)
        """
        self.task_name = task_name
        self.render_mode = render_mode
        self.max_episodes = max_episodes
        self.max_steps_per_episode = max_steps_per_episode
        self.seed = seed or np.random.randint(0, 1000000)
        self.use_tensorboard = use_tensorboard

        self.logger = logging.getLogger(__name__)
        self.env = None
        self.agent = None

        # Statistics tracking
        self.episode_rewards = []
        self.episode_lengths = []
        self.success_rate = 0.0

        # TensorBoard setup
        self.tb_writer = None
        if self.use_tensorboard:
            if log_dir is None:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                log_dir = f"runs/{self.task_name}_{timestamp}"

            os.makedirs(log_dir, exist_ok=True)
            self.tb_writer = SummaryWriter(log_dir)
            self.logger.info(f"TensorBoard logging enabled: {log_dir}")
            self.logger.info(f"View logs with: tensorboard --logdir {log_dir}")

    def setup_environment(self) -> gym.Env:
        """
        Set up the MetaWorld environment with MuJoCo rendering.

        Returns:
            Configured gymnasium environment
        """
        try:
            # Create MetaWorld environment
            if self.task_name == "reach-v3":
                # Use the reach task that matches our agent's policy
                mt1 = metaworld.MT1(self.task_name, seed=self.seed)
                env = mt1.train_classes[self.task_name]()
                task = mt1.train_tasks[0]
                env.set_task(task)
            else:
                # For other tasks, try to create them directly
                mt1 = metaworld.MT1(self.task_name, seed=self.seed)
                env = mt1.train_classes[self.task_name]()
                task = mt1.train_tasks[0]
                env.set_task(task)

            # Wrap with gymnasium if needed
            if not isinstance(env, gym.Env):
                env = gym.make(env.spec.id if hasattr(env, "spec") else self.task_name)

            # Configure rendering
            if hasattr(env, "render_mode"):
                env.render_mode = self.render_mode

            self.logger.info(f"Environment created: {self.task_name}")
            self.logger.info(f"Observation space: {env.observation_space}")
            self.logger.info(f"Action space: {env.action_space}")

            return env

        except Exception as e:
            self.logger.error(f"Failed to create environment {self.task_name}: {e}")
            self.logger.info("Falling back to reach-v3 environment")

            # Fallback to a simple reach environment
            mt1 = metaworld.MT1("reach-v3", seed=self.seed)
            env = mt1.train_classes["reach-v3"]()
            task = mt1.train_tasks[0]
            env.set_task(task)

            return env

    def setup_agent(self, env: gym.Env) -> RLAgent:
        """
        Set up the agent with the environment's observation and action spaces.

        Args:
            env: The gymnasium environment

        Returns:
            Configured RLAgent
        """
        agent = RLAgent(
            observation_space=env.observation_space,
            action_space=env.action_space,
            seed=self.seed,
            max_episode_steps=self.max_steps_per_episode,
        )

        self.logger.info("Agent initialized successfully")
        return agent

    def run_episode(self, episode_num: int) -> Dict[str, float]:
        """
        Run a single episode and return statistics.

        Args:
            episode_num: Episode number for logging

        Returns:
            Dictionary containing episode statistics
        """
        obs, info = self.env.reset(seed=self.seed + episode_num)
        self.agent.reset()

        episode_reward = 0.0
        episode_length = 0
        success = False
        step_rewards = []

        self.logger.info(f"Starting episode {episode_num + 1}")

        for step in range(self.max_steps_per_episode):
            try:
                # Get action from agent
                action_tensor = self.agent.act(obs)

                # Convert to numpy array if needed
                if hasattr(action_tensor, "numpy"):
                    action = action_tensor.numpy()
                elif hasattr(action_tensor, "detach"):
                    action = action_tensor.detach().numpy()
                else:
                    action = np.array(action_tensor)

                # Take step in environment
                obs, reward, terminated, truncated, info = self.env.step(action)

                # Render the environment for human viewing
                if self.render_mode == "human":
                    self.env.render()
                    time.sleep(0.02)  # Small delay to make visualization smoother

                episode_reward += reward
                episode_length += 1
                step_rewards.append(reward)

                # Log to TensorBoard (step-level metrics)
                if self.tb_writer:
                    global_step = episode_num * self.max_steps_per_episode + step
                    self.tb_writer.add_scalar("Step/Reward", reward, global_step)
                    self.tb_writer.add_scalar(
                        "Step/CumulativeReward", episode_reward, global_step
                    )

                # Check for success (MetaWorld specific)
                if hasattr(info, "get") and info.get("success", False):
                    success = True

                # Log progress occasionally
                if step % 50 == 0:
                    self.logger.debug(
                        f"Episode {episode_num + 1}, Step {step}: "
                        f"Reward {reward:.3f}, Total {episode_reward:.3f}"
                    )

                if terminated or truncated:
                    break

            except Exception as e:
                self.logger.error(f"Error during step {step}: {e}")
                break

        # Log episode-level metrics to TensorBoard
        if self.tb_writer:
            self.tb_writer.add_scalar("Episode/Reward", episode_reward, episode_num)
            self.tb_writer.add_scalar("Episode/Length", episode_length, episode_num)
            self.tb_writer.add_scalar("Episode/Success", float(success), episode_num)
            if step_rewards:
                self.tb_writer.add_scalar(
                    "Episode/AvgStepReward", np.mean(step_rewards), episode_num
                )
                self.tb_writer.add_scalar(
                    "Episode/MaxStepReward", np.max(step_rewards), episode_num
                )
                self.tb_writer.add_scalar(
                    "Episode/MinStepReward", np.min(step_rewards), episode_num
                )

        episode_stats = {
            "reward": episode_reward,
            "length": episode_length,
            "success": success,
        }

        self.logger.info(
            f"Episode {episode_num + 1} completed: "
            f"Reward {episode_reward:.3f}, "
            f"Length {episode_length}, "
            f"Success {success}"
        )

        return episode_stats

    def run_evaluation(self):
        """
        Run the complete evaluation session.
        """
        self.logger.info("Starting agent evaluation")

        # Setup environment and agent
        self.env = self.setup_environment()
        self.agent = self.setup_agent(self.env)

        # Run episodes
        total_successes = 0

        for episode in range(self.max_episodes):
            episode_stats = self.run_episode(episode)

            self.episode_rewards.append(episode_stats["reward"])
            self.episode_lengths.append(episode_stats["length"])

            if episode_stats["success"]:
                total_successes += 1

        # Calculate final statistics
        self.success_rate = total_successes / self.max_episodes
        avg_reward = np.mean(self.episode_rewards)
        avg_length = np.mean(self.episode_lengths)
        std_reward = np.std(self.episode_rewards)
        std_length = np.std(self.episode_lengths)

        # Log summary metrics to TensorBoard
        if self.tb_writer:
            self.tb_writer.add_scalar("Summary/AvgReward", avg_reward, 0)
            self.tb_writer.add_scalar("Summary/StdReward", std_reward, 0)
            self.tb_writer.add_scalar("Summary/AvgLength", avg_length, 0)
            self.tb_writer.add_scalar("Summary/StdLength", std_length, 0)
            self.tb_writer.add_scalar("Summary/SuccessRate", self.success_rate, 0)

            # Add histogram of rewards and lengths
            self.tb_writer.add_histogram(
                "Summary/RewardDistribution", np.array(self.episode_rewards), 0
            )
            self.tb_writer.add_histogram(
                "Summary/LengthDistribution", np.array(self.episode_lengths), 0
            )

            # Add hyperparameters
            self.tb_writer.add_hparams(
                {
                    "task": self.task_name,
                    "episodes": self.max_episodes,
                    "max_steps": self.max_steps_per_episode,
                    "seed": self.seed,
                    "render_mode": self.render_mode,
                },
                {
                    "avg_reward": avg_reward,
                    "success_rate": self.success_rate,
                    "avg_length": avg_length,
                },
            )

            self.tb_writer.flush()
            self.tb_writer.close()

        self.logger.info("=" * 50)
        self.logger.info("EVALUATION SUMMARY")
        self.logger.info("=" * 50)
        self.logger.info(f"Task: {self.task_name}")
        self.logger.info(f"Episodes: {self.max_episodes}")
        self.logger.info(f"Average Reward: {avg_reward:.3f} ± {std_reward:.3f}")
        self.logger.info(f"Average Length: {avg_length:.1f} ± {std_length:.1f}")
        self.logger.info(f"Success Rate: {self.success_rate:.1%}")
        if self.tb_writer:
            self.logger.info(
                "TensorBoard logs saved. View with: tensorboard --logdir runs/"
            )
        self.logger.info("=" * 50)

        # Close environment
        if self.env:
            self.env.close()

        return {
            "task": self.task_name,
            "episodes": self.max_episodes,
            "avg_reward": avg_reward,
            "std_reward": std_reward,
            "avg_length": avg_length,
            "std_length": std_length,
            "success_rate": self.success_rate,
            "episode_rewards": self.episode_rewards,
            "episode_lengths": self.episode_lengths,
        }

    def list_available_tasks(self):
        """
        List all available MetaWorld tasks.
        """
        try:
            # Get all MT1 tasks
            mt1_tasks = metaworld.MT1.get_train_tasks()
            self.logger.info("Available MetaWorld MT1 tasks:")
            for i, task in enumerate(mt1_tasks, 1):
                self.logger.info(f"  {i}. {task}")

            # Get all MT10 tasks
            mt10 = metaworld.MT10()
            self.logger.info("\nAvailable MetaWorld MT10 tasks:")
            for i, task in enumerate(mt10.train_classes.keys(), 1):
                self.logger.info(f"  {i}. {task}")

        except Exception as e:
            self.logger.error(f"Error listing tasks: {e}")
            self.logger.info("Some common MetaWorld tasks:")
            common_tasks = [
                "reach-v3",
                "push-v3",
                "pick-place-v3",
                "door-open-v3",
                "drawer-open-v3",
                "button-press-topdown-v3",
                "peg-insert-side-v3",
            ]
            for i, task in enumerate(common_tasks, 1):
                self.logger.info(f"  {i}. {task}")


def setup_logging(level=logging.INFO):
    """Configure logging for the evaluator."""
    logging.basicConfig(
        level=level,
        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
        handlers=[logging.StreamHandler(sys.stdout)],
    )


def main():
    """Main entry point for the evaluator."""
    parser = argparse.ArgumentParser(
        description="Evaluate the MetaWorld agent in MuJoCo"
    )
    parser.add_argument(
        "--task",
        type=str,
        default="reach-v3",
        help="MetaWorld task name (default: reach-v3)",
    )
    parser.add_argument(
        "--episodes",
        type=int,
        default=5,
        help="Number of episodes to run (default: 5)",
    )
    parser.add_argument(
        "--steps",
        type=int,
        default=200,
        help="Maximum steps per episode (default: 200)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed for reproducibility",
    )
    parser.add_argument(
        "--render-mode",
        type=str,
        default="human",
        choices=["human", "rgb_array"],
        help="Rendering mode (default: human)",
    )
    parser.add_argument(
        "--log-level",
        type=str,
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        help="Logging level (default: INFO)",
    )
    parser.add_argument(
        "--list-tasks",
        action="store_true",
        help="List available MetaWorld tasks and exit",
    )

    args = parser.parse_args()

    # Setup logging
    log_level = getattr(logging, args.log_level)
    setup_logging(log_level)

    # Create evaluator
    evaluator = AgentEvaluator(
        task_name=args.task,
        render_mode=args.render_mode,
        max_episodes=args.episodes,
        max_steps_per_episode=args.steps,
        seed=args.seed,
    )

    if args.list_tasks:
        evaluator.list_available_tasks()
        return

    try:
        evaluator.run_evaluation()
    except KeyboardInterrupt:
        logging.getLogger(__name__).info("Evaluation stopped by user")
    except Exception as e:
        logging.getLogger(__name__).error(
            f"Error during evaluation: {e}", exc_info=True
        )
        sys.exit(1)


if __name__ == "__main__":
    main()