create_model.py

import argparse
import json
import os
import shutil

from bridge.environment import FireboyWatergirlEnv
from profiles import (
    CURRENT_MODEL_GENERATION,
    DEFAULT_TARGET_LEVEL,
    ensure_profile_dirs,
    format_profile_label,
    get_profile_observation_mode,
    get_profile_paths,
    list_profile_names,
    normalize_observation_mode,
    OBSERVATION_MODES,
    write_profile_metadata,
)
from setting_catalog import PROFILE_SETTING_DEFAULTS


DEFAULT_SETTINGS = dict(PROFILE_SETTING_DEFAULTS)
DEFAULT_VISION_SETTINGS = {
    "observation_mode": FireboyWatergirlEnv.DEFAULT_OBSERVATION_MODE,
    "vision_width": FireboyWatergirlEnv.DEFAULT_VISION_WIDTH,
    "vision_height": FireboyWatergirlEnv.DEFAULT_VISION_HEIGHT,
    "vision_frame_stack": FireboyWatergirlEnv.DEFAULT_VISION_FRAME_STACK,
    "overlay_show_stacked_vision_frames": FireboyWatergirlEnv.DEFAULT_OVERLAY_SHOW_STACKED_VISION_FRAMES,
    "vision_crop_width": FireboyWatergirlEnv.DEFAULT_VISION_CROP_WIDTH,
    "vision_crop_height": FireboyWatergirlEnv.DEFAULT_VISION_CROP_HEIGHT,
    "policy": "MlpPolicy",
}


def _load_profile_config(profile):
    paths = ensure_profile_dirs(profile)
    if not os.path.exists(paths["config_path"]):
        return {}
    with open(paths["config_path"], "r", encoding="utf-8") as f:
        return json.load(f)


def _prompt_text(prompt, default=None):
    suffix = f" [{default}]" if default is not None else ""
    value = input(f"{prompt}{suffix}: ").strip()
    return value or default


def _prompt_bool(prompt, default):
    default_label = "y" if default else "n"
    while True:
        raw = _prompt_text(prompt, default_label).strip().lower()
        if raw in {"y", "yes", "true", "1"}:
            return True
        if raw in {"n", "no", "false", "0"}:
            return False
        print("Please enter y or n.")


def _prompt_model_type(default_mode):
    normalized_default = normalize_observation_mode(default_mode, "state")
    options = {
        "1": "state",
        "2": "vision",
        "3": "hybrid",
        "4": "vision_cropped",
        "5": "hybrid_cropped",
        "state": "state",
        "position": "state",
        "vision": "vision",
        "hybrid": "hybrid",
        "vision_cropped": "vision_cropped",
        "hybrid_cropped": "hybrid_cropped",
        "vision cropped": "vision_cropped",
        "hybrid cropped": "hybrid_cropped",
        "vc": "vision_cropped",
        "hc": "hybrid_cropped",
        "h": "hybrid",
        "v": "vision",
        "p": "state",
    }
    default_choice = {
        "state": "1",
        "vision": "2",
        "hybrid": "3",
        "vision_cropped": "4",
        "hybrid_cropped": "5",
    }[normalized_default]
    while True:
        print("Which model type do you want to create?")
        print("  1. Position model")
        print("  2. Vision model")
        print("  3. Hybrid model")
        print("  4. Vision cropped model")
        print("  5. Hybrid cropped model")
        raw = _prompt_text("Choose model type", default_choice).strip().lower()
        selected = options.get(raw)
        if selected:
            return selected
        print("Please choose 1, 2, 3, 4, or 5.")


def _choose_source_profile():
    profiles = [name for name in list_profile_names() if name != "default"]
    if not profiles:
        return None

    print("Existing profiles:")
    for index, profile in enumerate(profiles, start=1):
        print(f"  {index}. {format_profile_label(profile)}")

    choice = _prompt_text("Copy reward settings from which profile? Enter number or leave blank for defaults", "")
    if not choice:
        return None
    if choice.isdigit():
        index = int(choice) - 1
        if 0 <= index < len(profiles):
            return profiles[index]
    if choice in profiles:
        return choice
    print("Invalid selection. Using global defaults instead.")
    return None


def create_model(profile_name=None, copy_from=None, copy_weights=False, observation_mode=None):
    chosen_name = profile_name or _prompt_text("New model name", "variant4")
    target_paths = get_profile_paths(chosen_name)
    already_exists = os.path.isdir(target_paths["run_root"]) or os.path.exists(target_paths["swf_path"])
    if already_exists:
        print(f"Profile '{target_paths['profile']}' already exists.")
        return

    source_profile = copy_from
    if source_profile is None:
        source_profile = _choose_source_profile()

    source_config = _load_profile_config(source_profile) if source_profile else {}
    source_mode = get_profile_observation_mode(source_profile) if source_profile else FireboyWatergirlEnv.DEFAULT_OBSERVATION_MODE
    chosen_mode = (
        str(observation_mode).strip().lower()
        if observation_mode is not None
        else _prompt_model_type(source_mode)
    )
    chosen_mode = normalize_observation_mode(chosen_mode, "state")
    settings = dict(DEFAULT_SETTINGS)
    settings.update(DEFAULT_VISION_SETTINGS)
    for key in DEFAULT_SETTINGS:
        if key in source_config and source_config[key] is not None:
            settings[key] = source_config[key]
    for key in DEFAULT_VISION_SETTINGS:
        if key in source_config and source_config[key] is not None:
            settings[key] = source_config[key]
    settings["observation_mode"] = chosen_mode
    settings["policy"] = {
        "state": "MlpPolicy",
        "vision": "CnnPolicy",
        "hybrid": "MultiInputPolicy",
        "vision_cropped": "CnnPolicy",
        "hybrid_cropped": "MultiInputPolicy",
    }[chosen_mode]

    ensure_profile_dirs(chosen_name)
    paths = get_profile_paths(chosen_name)
    metadata = dict(settings)
    metadata["observation_config"] = FireboyWatergirlEnv.default_new_model_observation_config()
    metadata["model_generation"] = CURRENT_MODEL_GENERATION
    metadata["created_from"] = source_profile or "defaults"
    metadata["model_path"] = paths["model_path"]
    metadata["logs_dir"] = paths["logs_dir"]
    metadata["swf_path"] = paths["swf_path"]
    write_profile_metadata(chosen_name, metadata)

    if copy_weights and source_profile:
        source_paths = get_profile_paths(source_profile)
        source_model_zip = f"{source_paths['model_path']}.zip"
        target_model_zip = f"{paths['model_path']}.zip"
        if get_profile_observation_mode(source_profile) != chosen_mode:
            print(
                f"Skipped copying model weights because '{source_profile}' is "
                f"{get_profile_observation_mode(source_profile)} and '{paths['profile']}' is {chosen_mode}."
            )
        elif os.path.exists(source_model_zip):
            shutil.copy2(source_model_zip, target_model_zip)
            print(f"Copied model weights from '{source_profile}' to '{paths['profile']}'.")
        else:
            print(f"No saved weights found for '{source_profile}'. Created '{paths['profile']}' without weights.")

    type_label = {
        "state": "POSITION",
        "vision": "VISION",
        "hybrid": "HYBRID",
        "vision_cropped": "VISION CROPPED",
        "hybrid_cropped": "HYBRID CROPPED",
    }[chosen_mode]
    print(f"Created model profile '{paths['profile']}' [{type_label}].")
    print(f"SWF copy: {paths['swf_path']}")
    print(f"Config: {paths['config_path']}")
    print(f"Model path: {paths['model_path']}.zip")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--name", help="New model/profile name")
    parser.add_argument("--copy-from", help="Existing profile to copy reward settings from")
    parser.add_argument("--copy-weights", action="store_true", help="Also copy PPO weights from the source profile")
    parser.add_argument("--observation-mode", choices=sorted(OBSERVATION_MODES), help="Create a position/state, vision, hybrid, or cropped vision/hybrid profile")
    args = parser.parse_args()
    create_model(
        profile_name=args.name,
        copy_from=args.copy_from,
        copy_weights=args.copy_weights,
        observation_mode=args.observation_mode,
    )