CosmicClash/Game/scripts/ai_ship_controller.gd

class_name AIShipController
extends ShipController

# Drives a ship from a trained self-play policy (see TRAINING.md). Builds the
# same canonical observation as training (ShipObservations) and runs the
# policy MLP in GDScript (PolicyNetwork) — the shipped bot has no Python,
# .NET, or network dependency.
#
# Difficulty is (model, reaction_ticks, action_noise): weaker checkpoints make
# easier bots outright, and the two knobs handicap a given model further —
# slower reactions and noisier execution. Models live in res://bots/.

@export_file("*.json") var model_path: String = ""
# Decide a new action every N physics ticks, holding the last one between
# decisions. 8 matches the training action_repeat; larger = slower reactions.
@export_range(1, 60) var reaction_ticks: int = 8
# Uniform noise magnitude added to each action axis (0 = play at full skill).
@export_range(0.0, 1.0) var action_noise: float = 0.0

# Only meaningful for a "continuous"-action_space model (see
# ShipActionCodec) — i.e. one exported before curriculum generation 4, such
# as Game/bots/promoted/reference-grounded.json. Must mirror whatever the
# model was actually trained with: a model trained grounded (mask on) never
# got a reward gradient on these axes, so its raw output there is untrained
# noise — leaving this true for such a model doesn't make it fly well, it
# just lets that noise reach the ship instead of being discarded like it was
# in training. Set false to match a grounded-trained model's actual
# behaviour. Generation-4-onward (multi_discrete) models train the full
# action space from the start, so these flags are ignored for them.
@export var allow_vertical := true
@export var allow_pitch_roll := true

var _policy: PolicyNetwork
var _action := ShipAction.new()
var _ticks_until_decision := 0

# League mode revisits a small model pool every episode. Policies are
# immutable after load, so share parsed networks by path instead of reading
# and flattening the same 100KB+ JSON on every reset for every frozen ship.
static var _policy_cache: Dictionary = {}

var _ship: Ship
var _teammates: Array[Ship] = []
var _opponents: Array[Ship] = []
var _ball: RigidBody3D
var _attack_goal_position: Vector3
var _scene_refs_ready := false


func _ready():
	if not model_path.is_empty():
		load_policy(model_path)


# League training swaps a frozen opponent's policy between episodes without
# despawning the ship. Reset the held action/decision cadence with the model
# so no command from the previous opponent leaks into the next episode.
func load_policy(path: String) -> void:
	model_path = path
	if model_path.is_empty():
		_policy = null
	elif _policy_cache.has(model_path):
		_policy = _policy_cache[model_path]
	else:
		_policy = PolicyNetwork.load_from_file(model_path)
		if _policy != null:
			_policy_cache[model_path] = _policy
	_action = ShipAction.new()
	_ticks_until_decision = 0


func get_action() -> ShipAction:
	if _policy == null:
		return _action  # unloaded model: behaves like the inert placeholder
	if not _scene_refs_ready and not _discover_scene_refs():
		return _action

	_ticks_until_decision -= 1
	if _ticks_until_decision <= 0:
		_ticks_until_decision = reaction_ticks
		_decide()
	return _action


func _decide() -> void:
	var obs := ShipObservations.build(_ship, _teammates, _opponents, _ball, _attack_goal_position)
	var out := _policy.forward(obs)
	# See ShipActionCodec for the decode — the single source of truth shared
	# with the training side, so this must never reimplement layout/ordering
	# locally (see that file's header for why).
	if _policy.action_space.get("type", "continuous") == "continuous":
		_action = ShipActionCodec.from_continuous(out, action_noise)
		if not allow_pitch_roll:
			_action.rotation.x = 0.0
			_action.rotation.z = 0.0
		if not allow_vertical:
			_action.thrust.y = 0.0
	else:
		_action = ShipActionCodec.from_logits(out, action_noise)
	_action = ShipActionCodec.apply_team_frame(_action, _ship.team)


# Find ship/ball/teammates/opponents/goal once everything is spawned.
# ShipAction thrust axes are body-frame, while its rotation axes are mapped
# from the canonical team frame by ShipActionCodec. Rosters never change
# mid-match (no despawn path exists
# anywhere in this codebase), so this only needs to run once — sorted by
# spawn_index so a given ship keeps the same observation slot for the whole
# match, matching TrainingMode's identically-sorted lists.
func _discover_scene_refs() -> bool:
	_ship = get_parent() as Ship
	if _ship == null or not is_inside_tree():
		return false
	_ball = get_tree().get_first_node_in_group("ball")
	if _ball == null:
		return false
	# Cleared, not just appended to: if an earlier call reached this point but
	# a later check still failed, a retry must not re-append onto whatever it
	# already collected — that would duplicate every ship in the roster.
	_teammates.clear()
	_opponents.clear()
	for node in get_tree().get_nodes_in_group("ship"):
		var other := node as Ship
		if other == _ship:
			continue
		if other.team == _ship.team:
			_teammates.append(other)
		else:
			_opponents.append(other)
	_teammates.sort_custom(_by_spawn_index)
	_opponents.sort_custom(_by_spawn_index)
	for goal in get_tree().get_nodes_in_group("goal"):
		if goal.team == 1 - _ship.team:
			_attack_goal_position = goal.global_position
	_scene_refs_ready = true
	return true


static func _by_spawn_index(a: Ship, b: Ship) -> bool:
	return a.spawn_index < b.spawn_index