CosmicClash/training/test_action_space.py

"""Rung 0 of TRAINING.md's validation ladder: offline, no Godot, seconds to
run. Catches the single most likely silent killer in the generation-4
action-space redesign — a head-order mismatch between the Python trainer and
Game/scripts/ship_action_codec.gd's HEADS. If they disagree, training still
runs happily for 24h+ (pitch commands driving strafe thrusters, say) and
only surfaces as inexplicably-bad behaviour, not an error. This can't
directly parse the GDScript file, but it locks the two real, checkable
invariants an order mismatch would actually depend on: that gymnasium's Dict
key-sorting produces the exact order ship_action_codec.gd's HEADS is written
in, and that godot_rl's ActionSpaceProcessor converts that into the
MultiDiscrete nvec the trainer expects. export_policy.py's own index-level
parity check plus a real in-game round trip (rung 3) are what catch anything
this can't.

Usage:
    .venv/bin/python test_action_space.py
"""

import sys
from pathlib import Path

import gymnasium as gym
import numpy as np
from godot_rl.core.utils import ActionSpaceProcessor

from export_policy import ACTION_HEADS

# The order Game/scripts/ship_action_codec.gd's HEADS is written in — kept
# here as a literal, independent restatement (not derived from ACTION_HEADS)
# so this test can actually catch export_policy.py's own list being edited
# out of order too, not just catch nothing because both sides changed
# together.
EXPECTED_ORDER = ["rot_x", "rot_y", "rot_z", "thrust_x", "thrust_y", "thrust_z", "turbo"]
GAME_SCRIPTS = Path(__file__).resolve().parents[1] / "Game" / "scripts"


def check_action_heads_match_expected_order() -> None:
    names = [head["name"] for head in ACTION_HEADS]
    assert names == EXPECTED_ORDER, (
        f"export_policy.ACTION_HEADS order {names} != expected {EXPECTED_ORDER} — "
        "this must match Game/scripts/ship_action_codec.gd's HEADS exactly"
    )


def check_gymnasium_sorts_to_expected_order() -> None:
    # Build the Dict deliberately out of order (reversed) to prove it's
    # gymnasium's sort doing the work here, not insertion order — this is
    # exactly what godot_env.py does with the dict Godot sends over the wire
    # (see godot_env.py's from_dict, which builds a spaces.Dict from
    # ShipAIController.get_action_space()'s Dictionary).
    sizes = {head["name"]: len(head["bins"]) for head in ACTION_HEADS}
    reversed_dict = gym.spaces.Dict({name: gym.spaces.Discrete(sizes[name]) for name in reversed(EXPECTED_ORDER)})
    sorted_names = list(reversed_dict.keys())
    assert sorted_names == EXPECTED_ORDER, (
        f"gymnasium.spaces.Dict sorted {sorted_names}, expected {EXPECTED_ORDER} — "
        "if this changed, every export from this generation onward would be silently "
        "mis-ordered relative to ship_action_codec.gd"
    )


def check_action_space_processor_produces_expected_multi_discrete() -> None:
    sizes = [len(head["bins"]) for head in ACTION_HEADS]
    tuple_space = gym.spaces.Tuple([gym.spaces.Discrete(n) for n in sizes])
    processor = ActionSpaceProcessor(tuple_space, convert=True)
    assert isinstance(processor.action_space, gym.spaces.MultiDiscrete), (
        f"expected MultiDiscrete, got {type(processor.action_space)} — the all-discrete branch "
        "in godot_rl's ActionSpaceProcessor may have changed (see requirements.txt's pin note)"
    )
    assert list(processor.action_space.nvec) == sizes, (
        f"MultiDiscrete nvec {list(processor.action_space.nvec)} != expected {sizes}"
    )


def check_round_trip_preserves_per_head_values() -> None:
    # An integer action per env, one column per head in EXPECTED_ORDER —
    # confirms to_original_dist splits a MultiDiscrete action back into the
    # same per-head order it was built from (this is what set_action() on
    # the Godot side receives, keyed by head name).
    sizes = [len(head["bins"]) for head in ACTION_HEADS]
    tuple_space = gym.spaces.Tuple([gym.spaces.Discrete(n) for n in sizes])
    processor = ActionSpaceProcessor(tuple_space, convert=True)

    n_envs = 3
    rng = np.random.default_rng(0)
    action = np.stack([rng.integers(0, n, size=n_envs) for n in sizes], axis=1).astype(np.int64)
    original = processor.to_original_dist(action)
    assert len(original) == len(sizes)
    for head_index, expected_column in enumerate(action.T):
        np.testing.assert_array_equal(np.asarray(original[head_index]), expected_column)


def check_team_frame_mapping_is_shared() -> None:
    codec = (GAME_SCRIPTS / "ship_action_codec.gd").read_text()
    training = (GAME_SCRIPTS / "ship_ai_controller.gd").read_text()
    inference = (GAME_SCRIPTS / "ai_ship_controller.gd").read_text()
    assert "action.rotation.x = -action.rotation.x" in codec
    assert "action.rotation.z = -action.rotation.z" in codec
    assert "ShipActionCodec.apply_team_frame(" in training
    assert "ShipActionCodec.apply_team_frame(" in inference


def main() -> int:
    checks = [
        check_action_heads_match_expected_order,
        check_gymnasium_sorts_to_expected_order,
        check_action_space_processor_produces_expected_multi_discrete,
        check_round_trip_preserves_per_head_values,
        check_team_frame_mapping_is_shared,
    ]
    for check in checks:
        check()
        print(f"PASS: {check.__name__}")
    print(f"\nAll {len(checks)} action-space checks passed.")
    return 0


if __name__ == "__main__":
    sys.exit(main())