kaggle-environments 1.23.3__py3-none-any.whl → 1.23.4__py3-none-any.whl

kaggle_environments/envs/werewolf/scripts/plot_existing_trajectories.py

@@ -0,0 +1,135 @@
+import argparse
+import glob
+import json
+import logging
+import os
+import re
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from kaggle_environments.envs.werewolf.werewolf import CostSummary
+
+# Add the project root to the Python path to allow importing from kaggle_environments
+project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def plot_token_trajectories(trajectories_data, output_dir):
+    """
+    Plots token usage trajectories, grouped by max_turns, and saves them to files.
+    """
+    for metric, trajectories_by_turns in trajectories_data.items():
+        if not trajectories_by_turns:
+            logger.warning(f"No data found for metric '{metric}'. Skipping plot.")
+            continue
+
+        plt.figure(figsize=(12, 8))
+
+        # Create a color map for the different turn settings
+        turn_keys = sorted(trajectories_by_turns.keys(), key=int)
+        colors = plt.cm.viridis(np.linspace(0, 1, len(turn_keys)))
+        color_map = {turns: color for turns, color in zip(turn_keys, colors)}
+
+        for turns, trajectories in sorted(trajectories_by_turns.items(), key=lambda item: int(item[0])):
+            for i, traj in enumerate(trajectories):
+                if not all(isinstance(x, (int, float)) for x in traj):
+                    logger.error(
+                        f"Trajectory for metric '{metric}' (turns={turns}) contains non-numeric data. Skipping."
+                    )
+                    continue
+                # Only add a label to the first trajectory of each group for a clean legend
+                label = f"Max Turns: {turns}" if i == 0 else None
+                plt.plot(np.arange(len(traj)), traj, linestyle="-", alpha=0.4, color=color_map[turns], label=label)
+
+        plt.title(f"{metric.replace('_', ' ').title()} per Query Step Trajectories")
+        plt.xlabel("Query Step")
+        plt.ylabel(f"{metric.replace('_', ' ').title()} per Query Step")
+        plt.grid(True, which="both", linestyle="--", linewidth=0.5)
+        plt.legend()
+
+        plot_filename = os.path.join(output_dir, f"{metric}_trajectories.png")
+        plt.savefig(plot_filename)
+        plt.close()
+        logger.info(f"Saved trajectory plot: {plot_filename}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Load data from a measure_cost.py output directory and generate token trajectory plots."
+    )
+    parser.add_argument(
+        "-i",
+        "--input_dir",
+        type=str,
+        required=True,
+        help="Path to the output directory of a previous measure_cost.py run.",
+    )
+    args = parser.parse_args()
+
+    if not os.path.isdir(args.input_dir):
+        logger.error(f"Input directory not found: {args.input_dir}")
+        return
+
+    logger.info(f"Loading data from: {args.input_dir}")
+
+    all_trajectories = {"total_tokens": {}, "reasoning_tokens": {}, "text_tokens": {}}
+
+    # Find all game replay JSON files
+    game_files = glob.glob(os.path.join(args.input_dir, "game_*_run_*.json"))
+    if not game_files:
+        logger.error(f"No game replay files (game_*_run_*.json) found in {args.input_dir}.")
+        return
+
+    logger.info(f"Found {len(game_files)} game replay files to process.")
+
+    for game_file in game_files:
+        # Extract max_turns from filename
+        match = re.search(r"game_turns_(\d+)_run_", os.path.basename(game_file))
+        if not match:
+            logger.warning(f"Could not parse max_turns from filename: {game_file}. Skipping.")
+            continue
+        turns = match.group(1)
+
+        with open(game_file, "r") as f:
+            game_data = json.load(f)
+
+        cost_summary_dict = game_data.get("info", {}).get("GAME_END", {}).get("cost_summary")
+        if not cost_summary_dict:
+            logger.warning(f"No cost_summary found in {game_file}. Skipping.")
+            continue
+
+        cost_summary = CostSummary(**cost_summary_dict)
+
+        for agent_summary in cost_summary.cost_per_agent:
+            if agent_summary.data and agent_summary.data.usage_history:
+                usage_history_dicts = [usage.model_dump() for usage in agent_summary.data.usage_history]
+
+                total_tokens_traj = [usage.get("total_tokens", 0) or 0 for usage in usage_history_dicts]
+                all_trajectories["total_tokens"].setdefault(turns, []).append(total_tokens_traj)
+
+                reasoning_tokens_traj = [
+                    usage.get("completion_tokens_details", {}).get("reasoning_tokens", 0) or 0
+                    for usage in usage_history_dicts
+                ]
+                all_trajectories["reasoning_tokens"].setdefault(turns, []).append(reasoning_tokens_traj)
+
+                text_tokens_traj = [
+                    (u.get("completion_tokens", 0) or 0)
+                    - (u.get("completion_tokens_details", {}).get("reasoning_tokens", 0) or 0)
+                    for u in usage_history_dicts
+                ]
+                all_trajectories["text_tokens"].setdefault(turns, []).append(text_tokens_traj)
+
+    logger.info("Finished processing all files. Generating plots...")
+    plot_token_trajectories(all_trajectories, args.input_dir)
+    logger.info(f"--- Script finished. Plots saved in {args.input_dir} ---")
+
+
+if __name__ == "__main__":
+    main()

kaggle_environments/envs/werewolf/scripts/rerender_html.py

@@ -0,0 +1,87 @@
+import argparse
+import json
+import logging
+import os
+
+from kaggle_environments import make
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+
+def main():
+    """
+    Rerenders a Werewolf game replay HTML file from an existing game record JSON.
+    This is useful for updating the replay viewer to the latest version without
+    rerunning the entire game simulation.
+    """
+    parser = argparse.ArgumentParser(
+        description="Rerender a Werewolf game HTML replay from a JSON game record.",
+        formatter_class=argparse.RawTextHelpFormatter,
+    )
+    parser.add_argument(
+        "-i",
+        "--input_json",
+        type=str,
+        required=True,
+        help="Path to the input game record JSON file (e.g., werewolf_game.json).",
+    )
+    parser.add_argument(
+        "-o", "--output_html", type=str, required=True, help="Path to write the newly rendered HTML output file."
+    )
+    args = parser.parse_args()
+
+    logging.info(f"Loading game record from: {args.input_json}")
+    if not os.path.exists(args.input_json):
+        logging.error(f"Error: Input file not found at {args.input_json}")
+        return
+
+    try:
+        with open(args.input_json, "r", encoding="utf-8") as f:
+            replay_data = json.load(f)
+    except json.JSONDecodeError:
+        logging.error(f"Error: Failed to decode JSON from {args.input_json}. The file might be corrupted.")
+        return
+    except Exception as e:
+        logging.error(f"An unexpected error occurred while reading the file: {e}")
+        return
+
+    logging.info("Successfully loaded game data. Initializing Kaggle environment...")
+
+    # The environment name should be stored in the replay, but we default to 'werewolf'
+    env_name = replay_data.get("name", "werewolf")
+    if env_name != "werewolf":
+        logging.warning(f"Game record is for '{env_name}', but we are rendering with the 'werewolf' environment.")
+
+    try:
+        # Recreate the environment state from the replay file
+        env = make(
+            "werewolf",
+            configuration=replay_data.get("configuration"),
+            steps=replay_data.get("steps", []),
+            info=replay_data.get("info", {}),
+        )
+        logging.info("Environment initialized. Rendering new HTML...")
+
+        # Render the HTML. This will use the werewolf.js file included in the
+        # installed kaggle_environments package.
+        html_content = env.render(mode="html")
+
+        output_dir = os.path.dirname(args.output_html)
+        if output_dir:
+            os.makedirs(output_dir, exist_ok=True)
+
+        with open(args.output_html, "w", encoding="utf-8") as f:
+            f.write(html_content)
+
+        logging.info(f"Successfully rerendered HTML to: {args.output_html}")
+
+    except Exception as e:
+        logging.error(f"An error occurred during environment creation or rendering: {e}")
+        logging.error(
+            "Please ensure the 'kaggle_environments' package is correctly installed and the JSON file is valid."
+        )
+
+
+if __name__ == "__main__":
+    main()

kaggle_environments/envs/werewolf/scripts/run.py

@@ -0,0 +1,93 @@
+import argparse
+import logging
+import os
+import random
+
+import yaml
+
+from kaggle_environments.envs.werewolf.harness.base import LLMWerewolfAgent
+from kaggle_environments.envs.werewolf.runner import (
+    LogExecutionTime,
+    append_timestamp_to_dir,
+    log_git_hash,
+    run_werewolf,
+    setup_logger,
+)
+from kaggle_environments.envs.werewolf.werewolf import LLM_SYSTEM_PROMPT, AgentFactoryWrapper, register_agents
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run a single Werewolf game.")
+    parser.add_argument(
+        "-c",
+        "--config_path",
+        type=str,
+        default=os.path.join(os.path.dirname(__file__), "configs/run/run_config.yaml"),
+        help="Path to the YAML configuration file.",
+    )
+    parser.add_argument(
+        "-o", "--output_dir", type=str, default="werewolf_run", help="Output directory for the log and replay file."
+    )
+    parser.add_argument("-d", "--debug", action="store_true", help="Enable debug mode.")
+    parser.add_argument(
+        "-r", "--random_agents", action="store_true", help="Use random agents for all players for fast testing."
+    )
+    parser.add_argument(
+        "-a", "--append_timestamp_to_dir", action="store_true", help="Append a timestamp to the output directory."
+    )
+    parser.add_argument(
+        "-s", "--shuffle_roles", action="store_true", help="If provided, shuffle the roles provided in the config."
+    )
+
+    args = parser.parse_args()
+
+    # Create a unique subdirectory for this run
+    run_output_dir = append_timestamp_to_dir(args.output_dir, append=args.append_timestamp_to_dir)
+
+    os.makedirs(run_output_dir, exist_ok=True)
+
+    base_name = "werewolf_game"
+    setup_logger(output_dir=run_output_dir, base_name=base_name)
+
+    log_git_hash()
+
+    # Load game configuration
+    with open(args.config_path, "r") as f:
+        config = yaml.safe_load(f)
+        game_config = config.get("game_config", {})
+
+    # shuffle roles
+    if args.shuffle_roles:
+        role_and_params = [(agent["role"], agent.get("role_params", {})) for agent in game_config["agents"]]
+        random.shuffle(role_and_params)
+        for agent, (new_role, new_role_params) in zip(game_config["agents"], role_and_params):
+            agent["role"] = new_role
+            agent["role_params"] = new_role_params
+
+    # Extract agent harnesses from the config and register the agents
+    agents_ = [agent.get("agent_id", "random") for agent in game_config.get("agents", [])]
+    agent_dict = {}
+    for agent_name in agents_:
+        if agent_name.startswith("llm/"):
+            model_name = agent_name.lstrip("llm/")
+            agent_dict[agent_name] = AgentFactoryWrapper(
+                LLMWerewolfAgent, model_name=model_name, system_prompt=LLM_SYSTEM_PROMPT
+            )
+    register_agents(agent_dict)
+
+    if args.random_agents:
+        logger.info("Using random agents for all players.")
+        agents_ = ["random"] * len(agents_)
+
+    logger.info(f"Starting Werewolf game run. Output will be saved to: {run_output_dir}")
+    with LogExecutionTime(logger_obj=logger, task_str="single game"):
+        run_werewolf(
+            output_dir=run_output_dir, base_name=base_name, config=game_config, agents=agents_, debug=args.debug
+        )
+    logger.info(f"Game finished. Replay and log saved in: {run_output_dir}")
+
+
+if __name__ == "__main__":
+    main()

kaggle_environments/envs/werewolf/scripts/run_block.py

@@ -0,0 +1,237 @@
+import argparse
+import collections
+import logging
+import math
+import multiprocessing
+import os
+import random
+from itertools import permutations
+from typing import Any, Dict, List
+
+import tenacity
+import yaml
+from tqdm import tqdm
+
+from kaggle_environments.envs.werewolf.runner import LogExecutionTime, append_timestamp_to_dir, setup_logger
+from kaggle_environments.envs.werewolf.scripts.utils import run_single_game_cli
+
+# Initialize a placeholder logger
+logger = logging.getLogger(__name__)
+
+
+def load_config(config_path):
+    """Loads the configuration from a YAML file."""
+    with open(config_path, "r") as f:
+        return yaml.safe_load(f)
+
+
+def get_all_unique_role_configs(role_configs: List[Dict[str, Any]]) -> List[List[Dict[str, Any]]]:
+    """
+    Generates all unique permutations of role configurations.
+    A role configuration is a dict with 'role' and 'role_params'.
+    """
+
+    def make_hashable(config):
+        role = config["role"]
+        params = config.get("role_params", {})
+        if params:
+            return role, frozenset(params.items())
+        return role, frozenset()
+
+    def make_unhashable(hashable_config):
+        role, params_frozenset = hashable_config
+        return {"role": role, "role_params": dict(params_frozenset)}
+
+    hashable_configs = [make_hashable(c) for c in role_configs]
+    all_perms_hashable = list(set(permutations(hashable_configs)))
+    all_perms = [[make_unhashable(c) for c in p] for p in all_perms_hashable]
+    return all_perms
+
+
+run_single_game_with_retry = tenacity.retry(
+    wait=tenacity.wait_random_exponential(multiplier=1, min=2, max=10),
+    stop=tenacity.stop_after_attempt(3),
+    before_sleep=tenacity.before_sleep_log(logger, logging.INFO),
+)(run_single_game_cli)
+
+
+def game_runner_wrapper(args):
+    """Wrapper to unpack arguments for the multiprocessing pool."""
+    game_dir, game_config, use_random_agents, debug, _, _ = args
+    run_single_game_with_retry(game_dir, game_config, use_random_agents, debug)
+
+
+def generate_game_tasks(output_dir, num_blocks, config, use_random_agents, debug, shuffle_player_ids):
+    """
+    Generates all game configurations for the entire experiment.
+    """
+    base_game_config = config["game_config"]
+    players_data = base_game_config["agents"]
+    base_role_configs = [{"role": agent["role"], "role_params": agent.get("role_params", {})} for agent in players_data]
+
+    logger.info("Generating all unique role configurations...")
+    all_role_configs = get_all_unique_role_configs(base_role_configs)
+    logger.info(f"Found {len(all_role_configs)} unique arrangements.")
+
+    available_role_configs = []
+
+    for block_index in range(num_blocks):
+        block_dir = os.path.join(output_dir, f"block_{block_index}")
+        os.makedirs(block_dir, exist_ok=True)
+
+        if not available_role_configs:
+            if num_blocks > len(all_role_configs):
+                logger.warning("Sampling with replacement as num_blocks > unique configurations.")
+            available_role_configs = list(all_role_configs)
+            random.shuffle(available_role_configs)
+
+        block_role_config = available_role_configs.pop()
+        random.shuffle(players_data)
+        current_players_deque = collections.deque(players_data)
+
+        for game_in_block in range(len(players_data)):
+            game_dir = os.path.join(block_dir, f"game_{game_in_block}")
+            os.makedirs(game_dir, exist_ok=True)
+
+            current_players = list(current_players_deque)
+            game_agents_config = [
+                {**player_config, **block_role_config[i]} for i, player_config in enumerate(current_players)
+            ]
+
+            if shuffle_player_ids:
+                player_ids = [agent["id"] for agent in game_agents_config]
+                random.shuffle(player_ids)
+                for i, agent in enumerate(game_agents_config):
+                    agent["id"] = player_ids[i]
+
+            game_config = {**base_game_config, "agents": game_agents_config}
+            yield (game_dir, game_config, use_random_agents, debug, block_index, game_in_block)
+            current_players_deque.rotate(1)
+
+
+def run_experiment(
+    output_dir, num_blocks, config, use_random_agents, debug, parallel, num_processes, shuffle_player_ids
+):
+    """
+    Runs a tournament by generating all game tasks and processing them,
+    potentially in parallel.
+    """
+    if debug:
+        logger.warning("Debug mode is enabled. Forcing sequential execution.")
+
+    base_game_config = config["game_config"]
+    players_data = base_game_config["agents"]
+    total_games = num_blocks * len(players_data)
+
+    if parallel:
+        logger.info(f"Running games in parallel with up to {num_processes} processes.")
+
+    game_tasks = generate_game_tasks(output_dir, num_blocks, config, use_random_agents, debug, shuffle_player_ids)
+
+    with tqdm(total=total_games, desc="Processing Games") as pbar:
+        if parallel:
+            with multiprocessing.Pool(processes=num_processes) as pool:
+                for _ in pool.imap_unordered(game_runner_wrapper, game_tasks):
+                    pbar.update(1)
+        else:
+            for task_args in game_tasks:
+                game_runner_wrapper(task_args)
+                pbar.update(1)
+
+    logger.info("All game tasks have been processed.")
+
+
+def main():
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    default_config_path = os.path.join(script_dir, "configs", "run", "run_config.yaml")
+
+    parser = argparse.ArgumentParser(
+        description="Run a block-design experiment for the Werewolf game, "
+        "where each block is a complete role rotation amongst the players."
+    )
+    parser.add_argument(
+        "-o",
+        "--output_dir",
+        type=str,
+        help="Output directory for game replays and logs.",
+        default="werewolf_block_experiment",
+    )
+    parser.add_argument(
+        "-c", "--config", type=str, default=default_config_path, help="Path to the base configuration YAML file."
+    )
+    parser.add_argument(
+        "-b",
+        "--num_blocks",
+        type=int,
+        default=10,
+        help="Number of blocks to run. Each block is a complete role rotation.",
+    )
+    parser.add_argument(
+        "-r", "--use_random_agents", action="store_true", help="Use random agents for all players for fast testing."
+    )
+    parser.add_argument(
+        "-d",
+        "--debug",
+        action="store_true",
+        help="Enable debug mode for the game environment. "
+        "Note that you can use debug mode to enable intra game sequential execution.",
+    )
+    parser.add_argument("-p", "--parallel", action="store_true", help="Run games in parallel using multiple processes.")
+    parser.add_argument(
+        "-n", "--num_processes", type=int, default=None, help="Number of processes for parallel execution."
+    )
+    parser.add_argument(
+        "-a", "--append_timestamp_to_dir", action="store_true", help="Append a timestamp to the output directory."
+    )
+    parser.add_argument(
+        "-s",
+        "--shuffle_player_ids",
+        action="store_true",
+        help="Shuffle player ids for each game to account for name bias.",
+    )
+
+    args = parser.parse_args()
+
+    output_dir = append_timestamp_to_dir(args.output_dir, append=args.append_timestamp_to_dir)
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    setup_logger(output_dir, "run_block")
+
+    config = load_config(args.config)
+
+    num_players = len(config.get("game_config", {}).get("agents", []))
+    if args.num_processes is None:
+        num_processes = multiprocessing.cpu_count() * 0.9
+        if not args.debug:
+            num_processes /= num_players
+        num_processes = max(1, math.floor(num_processes))
+    else:
+        num_processes = args.num_processes
+
+    logger.info("Starting experiment with the following settings:")
+    logger.info(f"Output Directory: {output_dir}")
+    logger.info(f"Number of Blocks: {args.num_blocks}")
+    logger.info(f"Parallel Execution: {args.parallel}")
+    if args.parallel:
+        logger.info(f"Number of Processes: {num_processes}")
+    logger.info(f"Debug Mode: {args.debug}")
+    logger.info(f"Use Random Agents: {args.use_random_agents}")
+    logger.info(f"Shuffle Player IDs: {args.shuffle_player_ids}")
+
+    with LogExecutionTime(logger_obj=logger, task_str="block experiment"):
+        run_experiment(
+            output_dir=output_dir,
+            num_blocks=args.num_blocks,
+            config=config,
+            use_random_agents=args.use_random_agents,
+            debug=args.debug,
+            parallel=args.parallel,
+            num_processes=num_processes,
+            shuffle_player_ids=args.shuffle_player_ids,
+        )
+    logger.info("Experiment finished successfully.")
+
+
+if __name__ == "__main__":
+    main()