Upload submission from kinitro-agent-template

.gitignore

@@ -1,3 +1,5 @@
+runs/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[codz]
@@ -131,7 +133,7 @@ __pypackages__/
 celerybeat-schedule
 celerybeat.pid
 
-# Redis 
+# Redis
 *.rdb
 *.aof
 *.pid
@@ -195,9 +197,9 @@ cython_debug/
 .abstra/
 
 # Visual Studio Code
-#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
 #  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
-#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  and can be added to the global gitignore or merged into this file. However, if you prefer,
 #  you could uncomment the following to ignore the entire vscode folder
 # .vscode/
 
@@ -213,4 +215,4 @@ marimo/_lsp/
 __marimo__/
 
 # Streamlit
-.streamlit/secrets.toml
+.streamlit/secrets.toml

agent.capnp

@@ -1,13 +1,13 @@
-@0x893bac407c81b48c;
+@0xbf5147e1a2a3a3b1;
 
 interface Agent {
+  ping @0 (message :Text) -> (response :Text);
+  act  @1 (obs :Tensor) -> (action :Tensor);
+  reset @2 ();
+}
 
-    struct Tensor {
-        data @0   :Data; # tensor bytes tensor.numpy().tobytes()
-        shape @1  :List(UInt64); # tensor shape list(tensor.shape())
-        dtype @2  :Text; # data type name tensor.dtype()
-    }
-
-    act @0 (obs :Data) -> (action :Tensor);
-    reset @1 () -> ();
+struct Tensor {
+  data  @0 :Data;
+  shape @1 :List(Int32);
+  dtype @2 :Text;
 }

agent.py

@@ -8,6 +8,7 @@ import logging
 from typing import Any, Dict
 
 import gymnasium as gym
+
 import metaworld
 import numpy as np
 import torch

agent_server.py

@@ -1,15 +1,16 @@
 #!/usr/bin/env python3
 """
 Cap'n Proto RPC Server for Agent Interface
+Receives observation as Agent.Tensor (no pickle).
 """
 
 import asyncio
 import logging
 import os
-import pickle
+
+import capnp
 import numpy as np
 import torch
-import capnp
 
 # Load the schema
 schema_file = os.path.join(os.path.dirname(__file__), "agent.capnp")
@@ -17,6 +18,12 @@ agent_capnp = capnp.load(schema_file)
 
 logger = logging.getLogger(__name__)
 
+# Default network configuration
+DEFAULT_RPC_ADDRESS = "127.0.0.1"
+DEFAULT_RPC_PORT = 8000
+
+_TRAVERSAL_WORDS = 100 * 1024 * 1024  # match client; tune appropriately
+
 
 class AgentServer(agent_capnp.Agent.Server):
     """Cap'n Proto server implementation for AgentInterface"""
@@ -27,73 +34,80 @@ class AgentServer(agent_capnp.Agent.Server):
         self.logger.info("AgentServer initialized with agent: %s", type(agent).__name__)
 
     async def act(self, obs, **kwargs):
-        """Handle act RPC call"""
+        """Handle act RPC call. 'obs' is expected to be an Agent.Tensor struct."""
         try:
-            # Deserialize observation from bytes
-            observation = pickle.loads(obs)
-
-            # Call the agent's act method
-            action_tensor = self.agent.act(observation)
-
-            # Convert to numpy if it's a torch tensor
+            # obs is a struct with .data, .shape, .dtype
+            byte_len = len(obs.data) if obs and obs.data is not None else 0
+            self.logger.debug(
+                "Server.act invoked; incoming obs bytes=%d shape=%s dtype=%s",
+                byte_len,
+                list(obs.shape) if obs else None,
+                obs.dtype if obs else None,
+            )
+
+            # reconstruct numpy observation
+            obs_np = np.frombuffer(obs.data, dtype=np.dtype(obs.dtype)).reshape(
+                tuple(obs.shape)
+            )
+
+            # call the underlying agent synchronously (user's agent.act should accept ndarray)
+            action_tensor = self.agent.act(obs_np)
+
+            # convert to numpy
             if isinstance(action_tensor, torch.Tensor):
-                action_numpy = action_tensor.detach().cpu().numpy()
+                action_np = action_tensor.detach().cpu().numpy()
             else:
-                action_numpy = np.array(action_tensor)
-
-            # Prepare tensor response
-            response = agent_capnp.Agent.Tensor.new_message()
-            response.data = action_numpy.tobytes()
-            response.shape = list(action_numpy.shape)
-            response.dtype = str(action_numpy.dtype)
+                action_np = np.array(action_tensor)
 
+            # Build response Tensor
+            response = agent_capnp.Tensor.new_message()
+            response.data = action_np.tobytes()
+            response.shape = [int(s) for s in action_np.shape]
+            response.dtype = str(action_np.dtype)
             return response
-        except Exception as e:
-            self.logger.error(f"Error in act: {e}", exc_info=True)
+        except Exception:
+            self.logger.exception("Exception in AgentServer.act")
             raise
 
     async def reset(self, **kwargs):
-        """Handle reset RPC call"""
         try:
             self.agent.reset()
-        except Exception as e:
-            self.logger.error(f"Error in reset: {e}", exc_info=True)
+        except Exception:
+            self.logger.exception("Error in reset")
             raise
 
+    async def ping(self, message, **kwargs):
+        self.logger.info(f"Ping received: {message}")
+        return "pong"
+
 
-async def serve(agent, address="127.0.0.1", port=8000):
+async def serve(agent, address=DEFAULT_RPC_ADDRESS, port=DEFAULT_RPC_PORT):
     """Serve the agent using asyncio approach"""
 
     async def new_connection(stream):
-        """Handler for each new client connection"""
         try:
-            # Create TwoPartyServer for this connection
-            server = capnp.TwoPartyServer(stream, bootstrap=AgentServer(agent))
-
-            # Wait for the connection to disconnect
+            server = capnp.TwoPartyServer(
+                stream,
+                bootstrap=AgentServer(agent),
+                traversal_limit_in_words=_TRAVERSAL_WORDS,
+            )
             await server.on_disconnect()
+        except Exception:
+            logger.exception("Error handling connection")
 
-        except Exception as e:
-            logger.error(f"Error handling connection: {e}", exc_info=True)
-
-    # Create the server
     server = await capnp.AsyncIoStream.create_server(new_connection, address, port)
-
-    logger.info(f"Agent RPC server listening on {address}:{port}")
+    logger.info("Agent RPC server listening on %s:%d", address, port)
 
     try:
-        # Keep the server running
         async with server:
             await server.serve_forever()
-    except Exception as e:
-        logger.error(f"Server error: {e}", exc_info=True)
+    except Exception:
+        logger.exception("Server error")
     finally:
         logger.info("Server shutting down")
 
 
-def start_server(agent, address="127.0.0.1", port=8000):
-    """Start server with proper asyncio event loop handling"""
-
+def start_server(agent, address=DEFAULT_RPC_ADDRESS, port=DEFAULT_RPC_PORT):
     async def run_server_with_kj():
         async with capnp.kj_loop():
             await serve(agent, address, port)
@@ -104,9 +118,7 @@ def start_server(agent, address="127.0.0.1", port=8000):
         logger.info("Server stopped by user")
 
 
-def run_server_in_process(agent, address="127.0.0.1", port=8000):
-    """Entry point for running server in a separate process"""
-
+def run_server_in_process(agent, address=DEFAULT_RPC_ADDRESS, port=DEFAULT_RPC_PORT):
     async def run_with_kj():
         async with capnp.kj_loop():
             await serve(agent, address, port)

evaluation.py

@@ -0,0 +1,499 @@
+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()

main.py

@@ -1,17 +1,22 @@
 #!/usr/bin/env python3
 """
-Main entry point for the agent server.
+Main entry point for the agent server and evaluation.
 
-This script creates an agent implementation and starts the RPC server
-to handle requests from the evaluator.
+This script provides multiple commands:
+- server: Creates an agent implementation and starts the RPC server
+- eval: Runs local evaluation of the agent with visual rendering
 """
 
 import argparse
 import logging
+import subprocess
 import sys
+import threading
+import time
+import webbrowser
 
 from agent import RLAgent
-from agent_server import start_server
+from evaluation import AgentEvaluator
 
 
 def setup_logging(level=logging.INFO):
@@ -23,16 +28,81 @@ def setup_logging(level=logging.INFO):
     )
 
 
+def launch_tensorboard(log_dir, port=6006):
+    """Launch TensorBoard in a separate thread."""
+
+    def run_tensorboard():
+        try:
+            # Wait a moment for initial logs to be written
+            time.sleep(2)
+
+            # Launch TensorBoard
+            subprocess.run(
+                [
+                    "tensorboard",
+                    "--logdir",
+                    log_dir,
+                    "--port",
+                    str(port),
+                    "--host",
+                    "localhost",
+                    "--reload_interval",
+                    "1",
+                ],
+                check=True,
+                capture_output=True,
+            )
+        except subprocess.CalledProcessError:
+            # TensorBoard failed to start, but don't crash the evaluation
+            pass
+        except FileNotFoundError:
+            # TensorBoard not installed
+            pass
+
+    # Start TensorBoard in background thread
+    tb_thread = threading.Thread(target=run_tensorboard, daemon=True)
+    tb_thread.start()
+
+    # Give TensorBoard a moment to start
+    time.sleep(3)
+
+    # Try to open browser
+    try:
+        webbrowser.open(f"http://localhost:{port}")
+    except Exception:
+        # Browser opening failed, but that's okay
+        pass
+
+    return f"http://localhost:{port}"
+
+
 def main():
     """Main entry point."""
-    parser = argparse.ArgumentParser(description="Start the agent server")
-    parser.add_argument(
-        "--host", type=str, default="*", help="Host to bind the server to"
+    parser = argparse.ArgumentParser(
+        description="Agent server and evaluation tool",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python main.py server --host localhost --port 8000
+  python main.py eval --task reach-v3 --episodes 5
+  python main.py eval --task push-v3 --episodes 10 --render-mode rgb_array
+  python main.py eval --task reach-v3 --episodes 20 --no-tensorboard
+  python main.py eval --task door-open-v3 --log-dir custom_logs/
+        """,
     )
-    parser.add_argument(
+
+    # Add subcommands
+    subparsers = parser.add_subparsers(dest="command", help="Available commands")
+
+    # Server subcommand
+    server_parser = subparsers.add_parser("server", help="Start the agent server")
+    server_parser.add_argument(
+        "--host", type=str, default="0.0.0.0", help="Host to bind the server to"
+    )
+    server_parser.add_argument(
         "--port", type=int, default=8000, help="Port to bind the server to"
     )
-    parser.add_argument(
+    server_parser.add_argument(
         "--log-level",
         type=str,
         default="INFO",
@@ -40,13 +110,97 @@ def main():
         help="Logging level",
     )
 
+    # Evaluation subcommand
+    eval_parser = subparsers.add_parser("eval", help="Run local agent evaluation")
+    eval_parser.add_argument(
+        "--task",
+        type=str,
+        default="reach-v3",
+        help="MetaWorld task name (default: reach-v3)",
+    )
+    eval_parser.add_argument(
+        "--episodes",
+        type=int,
+        default=5,
+        help="Number of episodes to run (default: 5)",
+    )
+    eval_parser.add_argument(
+        "--steps",
+        type=int,
+        default=200,
+        help="Maximum steps per episode (default: 200)",
+    )
+    eval_parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Random seed for reproducibility",
+    )
+    eval_parser.add_argument(
+        "--render-mode",
+        type=str,
+        default="human",
+        choices=["human", "rgb_array"],
+        help="Rendering mode (default: human)",
+    )
+    eval_parser.add_argument(
+        "--log-level",
+        type=str,
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Logging level (default: INFO)",
+    )
+    eval_parser.add_argument(
+        "--list-tasks",
+        action="store_true",
+        help="List available MetaWorld tasks and exit",
+    )
+    eval_parser.add_argument(
+        "--tensorboard",
+        action="store_true",
+        default=True,
+        help="Enable TensorBoard logging (default: True)",
+    )
+    eval_parser.add_argument(
+        "--no-tensorboard",
+        action="store_true",
+        help="Disable TensorBoard logging",
+    )
+    eval_parser.add_argument(
+        "--log-dir",
+        type=str,
+        default=None,
+        help="TensorBoard log directory (auto-generated if not specified)",
+    )
+
     args = parser.parse_args()
 
+    # If no command is provided, show help
+    if not args.command:
+        parser.print_help()
+        sys.exit(1)
+
     # Setup logging
     log_level = getattr(logging, args.log_level)
     setup_logging(log_level)
     logger = logging.getLogger(__name__)
 
+    if args.command == "server":
+        run_server(args, logger)
+    elif args.command == "eval":
+        run_evaluation(args, logger)
+
+
+def run_server(args, logger):
+    """Run the agent server."""
+    # Import server functionality only when needed to avoid capnp dependency for eval
+    try:
+        from agent_server import start_server
+    except ImportError as e:
+        logger.error(f"Failed to import server functionality: {e}")
+        logger.error("Make sure capnp and other server dependencies are installed")
+        sys.exit(1)
+
     logger.info(f"Starting agent server on {args.host}:{args.port}")
 
     # Create the RLAgent
@@ -62,5 +216,72 @@ def main():
         sys.exit(1)
 
 
+def run_evaluation(args, logger):
+    """Run local agent evaluation."""
+    logger.info("Running local evaluation")
+
+    # Determine TensorBoard usage
+    use_tensorboard = args.tensorboard and not args.no_tensorboard
+
+    # Setup log directory if using TensorBoard
+    log_dir = args.log_dir
+    if use_tensorboard and not log_dir:
+        from datetime import datetime
+
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        log_dir = f"runs/{args.task}_{timestamp}"
+
+    # 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,
+        use_tensorboard=use_tensorboard,
+        log_dir=log_dir,
+    )
+
+    if args.list_tasks:
+        evaluator.list_available_tasks()
+        return
+
+    # Launch TensorBoard if enabled
+    tensorboard_url = None
+    if use_tensorboard and log_dir:
+        logger.info("Starting TensorBoard...")
+        try:
+            tensorboard_url = launch_tensorboard(log_dir)
+            logger.info(f"TensorBoard available at: {tensorboard_url}")
+            logger.info("TensorBoard will show metrics in real-time during evaluation")
+        except Exception as e:
+            logger.warning(f"Failed to start TensorBoard: {e}")
+            logger.info("Continuing evaluation without TensorBoard...")
+
+    try:
+        evaluator.run_evaluation()
+        logger.info("Evaluation completed successfully")
+
+        if tensorboard_url:
+            logger.info(f"View detailed metrics at: {tensorboard_url}")
+            logger.info("TensorBoard will continue running in the background")
+
+        # Optionally save results to file
+        # import json
+        # with open("evaluation_results.json", "w") as f:
+        #     json.dump(results, f, indent=2)
+        # logger.info("Results saved to evaluation_results.json")
+
+    except KeyboardInterrupt:
+        logger.info("Evaluation stopped by user")
+        if tensorboard_url:
+            logger.info(f"TensorBoard still available at: {tensorboard_url}")
+    except Exception as e:
+        logger.error(f"Error during evaluation: {e}", exc_info=True)
+        if tensorboard_url:
+            logger.info(f"TensorBoard still available at: {tensorboard_url}")
+        sys.exit(1)
+
+
 if __name__ == "__main__":
     main()

pyproject.toml

@@ -6,7 +6,13 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
     "metaworld>=3.0.0",
-    "torch>=2.8.0"
+    "torch>=2.8.0",
+    "gymnasium>=0.29.0",
+    "mujoco>=3.0.0",
+    "numpy>=1.24.0",
+    "pycapnp>=2.1.0",
+    "tensorboard>=2.15.0",
+    "matplotlib>=3.7.0"
 ]
 
 [dependency-groups]