CosmicClash/training/cosmic_env.py

"""Godot RL Agents environment wrappers that run Cosmic Clash from source.

Stock GodotEnv expects an *exported* game executable and rewrites its path
per-platform. These subclasses launch the project straight from the repo with
a Godot binary instead (no export step), pointing it at the training scene.
Each Godot instance contributes two agents (one ship per team) that share the
learning policy: self-play by construction.
"""

import pathlib
import subprocess

from godot_rl.core.godot_env import GodotEnv
from godot_rl.wrappers.stable_baselines_wrapper import StableBaselinesGodotEnv

REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent
GAME_DIR = REPO_ROOT / "Game"
TRAINING_SCENE = "res://scenes/training.tscn"


class CosmicClashEnv(GodotEnv):
    """GodotEnv that launches `godot --path Game res://scenes/training.tscn`."""

    # env_path is a Godot binary, not an exported game: skip the suffix and
    # platform checks stock GodotEnv applies to exported executables.
    def _set_platform_suffix(self, env_path: str) -> str:
        return env_path

    def check_platform(self, filename: str):
        pass

    def _launch_env(self, env_path, port, show_window, framerate, seed, action_repeat, speedup, **kwargs):
        # sync.gd reads --key=value pairs from the raw command line; they must
        # NOT go after a `--` separator or OS.get_cmdline_args() drops them.
        cmd = [
            env_path,
            "--path",
            str(GAME_DIR),
            TRAINING_SCENE,
            f"--port={port}",
            f"--env_seed={seed}",
        ]
        if not show_window:
            cmd += ["--headless", "--disable-render-loop"]
        if framerate is not None:
            cmd += ["--fixed-fps", str(framerate)]
        if action_repeat is not None:
            cmd.append(f"--action_repeat={action_repeat}")
        if speedup is not None:
            cmd.append(f"--speedup={speedup}")
        for key, value in kwargs.items():
            cmd.append(f"--{key}={value}")
        self.proc = subprocess.Popen(cmd, start_new_session=True)


class CosmicClashVecEnv(StableBaselinesGodotEnv):
    """SB3 VecEnv over N parallel CosmicClashEnv instances.

    convert_action_space=True flattens the env's (Box(6), Discrete(2)) action
    space into a single Box(7): thrust xyz, rotation xyz, turbo (>0 means on).
    """

    def __init__(self, godot_bin: str, n_parallel: int = 1, seed: int = 0, port: int = GodotEnv.DEFAULT_PORT, **kwargs):
        self.envs = [
            CosmicClashEnv(
                env_path=godot_bin,
                convert_action_space=True,
                port=port + p,
                seed=seed + p,
                **kwargs,
            )
            for p in range(n_parallel)
        ]
        self.n_parallel = n_parallel
        self._check_valid_action_space()
        self.results = None