CosmicClash/Game/scripts/training_mode.gd

class_name TrainingMode
extends GameMode

# Headless self-play training mode: two RL-driven ships, no HUD, no camera.
# The scene also contains the godot_rl_agents Sync node, which speaks TCP to
# the Python trainer; this mode owns the environment rules — episodes, goal
# rewards, and randomized episode-start states (the RLGym "state setter"
# lesson: varied starts massively speed up learning versus kickoff-only).
#
# Run: godot --headless --path Game res://scenes/training.tscn
# (started automatically by training/train.py; boots into idle ships with a
# warning if no trainer is listening).
#
# Eval mode (used by training/evaluate.py): pass --eval_model_a=<path> and
# --eval_model_b=<path> (+ optional --eval_episodes=N) and both ships are
# instead driven by those exported policies via AIShipController; each episode
# ends at the first goal (or a draw on timeout), and a final machine-readable
# "EVAL_RESULT {...}" line is printed before quitting.

@export var episode_length_seconds := 30.0
# Run01-vs-run02 eval (see training/eval_history.json) came back 87.5% draws:
# with a 30s episode, the dense per-tick terms on ShipAIController can sum to
# several times this value before it was raised, so scoring and forfeiting
# the rest of the episode's farmable reward was worse than never finishing.
# Raised well above that ceiling so a real scoring chance always beats
# continuing to farm dense reward for however long is left in the episode.
@export var goal_reward := 40.0

# Episode-start state mix; remaining probability = fully random state.
@export_range(0.0, 1.0) var kickoff_state_chance := 0.2
# Raised from 0.2: fixing the reward incentive to score (see goal_reward,
# ball_touch_reward, time_penalty on ShipAIController) only helps if the
# policy also gets enough reps at actually finishing. At 0.2 that scenario
# was 1 in 5 episode starts; most training time was spent in generic
# midfield play where a finish never comes up.
@export_range(0.0, 1.0) var ball_near_goal_chance := 0.35

# Placement bounds for randomized episode starts, derived from the standard
# enclosure (ArenaBoundary). The inset keeps a randomly oriented ship (1x1x4
# box, worst-case half-extent ~2.05) from spawning intersecting the walls,
# ceiling, or goal sensors.
const SPAWN_INSET := 2.5
const FIELD_HALF_X := ArenaBoundary.INNER_HALF_X - SPAWN_INSET
const FIELD_HALF_Z := ArenaBoundary.GOAL_LINE_Z - SPAWN_INSET
const FIELD_MIN_Y := 1.5
const FIELD_MAX_Y := ArenaBoundary.INNER_HEIGHT - SPAWN_INSET
# The corner curves reach at most their chord plane |x| + |z| = INNER_HALF_X
# + INNER_HALF_Z - CORNER_RADIUS; spawns keep the same SPAWN_INSET clearance
# from that plane as from the walls (perpendicular distance, hence the
# sqrt(2) when expressed in |x| + |z| terms). The true curve bulges outward
# from the chord, so this is conservative.
const CORNER_LIMIT := ArenaBoundary.INNER_HALF_X + ArenaBoundary.INNER_HALF_Z \
	- ArenaBoundary.CORNER_RADIUS - SPAWN_INSET * sqrt(2.0)
# Below this height a tilted ship could reach down into the wall-base
# fillets, so low spawns stay an extra BASE_RADIUS off the walls.
const FILLET_CLEAR_Y := ArenaBoundary.BASE_RADIUS + FIELD_MIN_Y
const MAX_RANDOM_BALL_SPEED := 12.0
const MAX_RANDOM_SHIP_SPEED := 8.0

# Sim runs at 60 physics ticks per sim-second regardless of speedup.
const TICKS_PER_SIM_SECOND := 60.0

# The arena is physically enclosed, so nothing should ever get this far out.
# If a body escapes anyway (physics regression, boundary edit), it is warned
# about and respawned with no reward change and no episode end — a multi-hour
# training run must survive it, and the escape must not shape rewards.
const ESCAPE_MARGIN := 15.0

var _agents: Array[ShipAIController] = []

# Eval mode state (see header comment)
var _eval := false
var _eval_models: Array[String] = ["", ""]
var _eval_episodes := 20
var _eval_goals := {0: 0, 1: 0}
var _eval_draws := 0
var _eval_episodes_done := 0
var _episode_ticks := 0


func _start() -> void:
	_parse_eval_args()
	spawn_ball()
	if _eval:
		for team in [0, 1]:
			var bot := AIShipController.new()
			bot.model_path = _eval_models[team]
			spawn_ship(team, 0, bot)
		return
	var ship_team0 := spawn_ship(0, 0, RLShipController.new())
	var ship_team1 := spawn_ship(1, 0, RLShipController.new())
	_attach_agent(ship_team0, ship_team1)
	_attach_agent(ship_team1, ship_team0)


func _parse_eval_args() -> void:
	var args := {}
	for argument in OS.get_cmdline_args():
		if argument.begins_with("--") and argument.find("=") > -1:
			var key_value := argument.lstrip("--").split("=", true, 1)
			args[key_value[0]] = key_value[1]
	if args.has("eval_model_a") and args.has("eval_model_b"):
		_eval = true
		_eval_models[0] = args["eval_model_a"]
		_eval_models[1] = args["eval_model_b"]
		_eval_episodes = int(args.get("eval_episodes", str(_eval_episodes)))


func _attach_agent(ship: Ship, opponent: Ship) -> void:
	var agent := ShipAIController.new()
	agent.name = "ShipAIController"
	agent.reset_after = int(episode_length_seconds * TICKS_PER_SIM_SECOND)
	ship.add_child(agent)
	agent.setup(ship, ship.controller as RLShipController, ball, opponent, _attack_goal_position(ship.team))
	_agents.append(agent)


# The goal this team scores into: the one the opponent defends/concedes.
func _attack_goal_position(team: int) -> Vector3:
	for goal in arena.get_goals():
		if goal.team == 1 - team:
			return goal.global_position
	push_error("TrainingMode: no goal found for team %d to attack" % team)
	return Vector3.ZERO


func _physics_process(_delta):
	_respawn_escaped_bodies()
	if _eval:
		_episode_ticks += 1
		if _episode_ticks > int(episode_length_seconds * TICKS_PER_SIM_SECOND):
			_eval_draws += 1
			_end_eval_episode()
		return

	# Both trainer-requested resets and truncation (reset_after ticks elapsed)
	# surface as needs_reset. Only truncation is an episode end the trainer
	# must be told about via done — a trainer-requested reset already knows.
	var needs_reset := false
	var truncated := false
	for agent in _agents:
		needs_reset = needs_reset or agent.needs_reset
		truncated = truncated or agent.n_steps > agent.reset_after
	if needs_reset:
		if truncated:
			for agent in _agents:
				agent.done = true
		_reset_episode()
		return


# Escape failsafe: see ESCAPE_MARGIN.
func _respawn_escaped_bodies() -> void:
	for ship in ships:
		if is_instance_valid(ship) and _escaped(ship.global_position):
			push_warning("TrainingMode: ship escaped the enclosed arena — check boundary colliders")
			_place_body(ship, _ship_spawn_transforms[ship], Vector3.ZERO, Vector3.ZERO)
	if is_instance_valid(ball) and _escaped(ball.global_position):
		push_warning("TrainingMode: ball escaped the enclosed arena — check boundary colliders")
		_place_body(ball, arena.get_ball_spawn(), Vector3.ZERO, Vector3.ZERO)


func _escaped(position: Vector3) -> bool:
	return absf(position.x) > ArenaBoundary.INNER_HALF_X + ESCAPE_MARGIN \
		or absf(position.z) > ArenaBoundary.INNER_HALF_Z + ESCAPE_MARGIN \
		or position.y < -ESCAPE_MARGIN \
		or position.y > ArenaBoundary.INNER_HEIGHT + ESCAPE_MARGIN


func _on_goal_scored(conceding_team: int) -> void:
	if _eval:
		_eval_goals[1 - conceding_team] += 1
		_end_eval_episode()
		return
	for agent in _agents:
		agent.reward += goal_reward if agent.ship.team != conceding_team else -goal_reward
		agent.done = true
	_reset_episode()


func _end_eval_episode() -> void:
	_eval_episodes_done += 1
	_episode_ticks = 0
	if _eval_episodes_done >= _eval_episodes:
		print("EVAL_RESULT " + JSON.stringify({
			"model_a": _eval_models[0],
			"model_b": _eval_models[1],
			"episodes": _eval_episodes_done,
			"goals_a": _eval_goals[0],
			"goals_b": _eval_goals[1],
			"draws": _eval_draws,
		}))
		get_tree().quit()
		return
	# Randomized states (not kickoff): deterministic policies would otherwise
	# replay the identical episode every time.
	_reset_episode()


func _reset_episode() -> void:
	for agent in _agents:
		agent.reset()

	var roll := randf()
	if roll < kickoff_state_chance:
		reset_ball()
		reset_ships()
	elif roll < kickoff_state_chance + ball_near_goal_chance:
		_place_ships_random()
		_place_ball_near_goal()
	else:
		_place_ships_random()
		_place_ball_random()


func _place_ball_random() -> void:
	var velocity := _random_direction() * randf_range(0.0, MAX_RANDOM_BALL_SPEED)
	_place_body(ball, Transform3D(Basis.IDENTITY, _random_position()), velocity, Vector3.ZERO)


# Attacking/defending drill states: ball close to a goal, moving toward it.
func _place_ball_near_goal() -> void:
	var goals := arena.get_goals()
	var goal: Goal = goals[randi() % goals.size()]
	var toward_centre := -signf(goal.global_position.z)
	var position := Vector3(
		randf_range(-4.0, 4.0),
		randf_range(FIELD_MIN_Y, 4.0),
		goal.global_position.z + toward_centre * randf_range(3.0, 6.0)
	)
	var to_goal := (goal.global_position - position).normalized()
	var velocity := (to_goal + _random_direction() * 0.3).normalized() * randf_range(2.0, MAX_RANDOM_BALL_SPEED)
	_place_body(ball, Transform3D(Basis.IDENTITY, position), velocity, Vector3.ZERO)


func _place_ships_random() -> void:
	for ship in ships:
		var orientation := Basis.from_euler(Vector3(
			randf_range(-0.4, 0.4),
			randf_range(-PI, PI),
			randf_range(-0.4, 0.4)
		))
		var velocity := _random_direction() * randf_range(0.0, MAX_RANDOM_SHIP_SPEED)
		_place_body(ship, Transform3D(orientation, _random_position()), velocity, Vector3.ZERO)


func _random_position() -> Vector3:
	# Resample anything too close to a corner curve or wall-base fillet (see
	# CORNER_LIMIT / FILLET_CLEAR_Y); the violating region is a few percent
	# of the volume, so 20 attempts effectively never fall through.
	var position := Vector3.ZERO
	for _attempt in 20:
		position = Vector3(
			randf_range(-FIELD_HALF_X, FIELD_HALF_X),
			randf_range(FIELD_MIN_Y, FIELD_MAX_Y),
			randf_range(-FIELD_HALF_Z, FIELD_HALF_Z)
		)
		if _spawn_position_clear(position):
			break
	return position


func _spawn_position_clear(position: Vector3) -> bool:
	if absf(position.x) + absf(position.z) > CORNER_LIMIT:
		return false
	if position.y >= FILLET_CLEAR_Y:
		return true
	return absf(position.x) <= FIELD_HALF_X - ArenaBoundary.BASE_RADIUS \
		and absf(position.z) <= FIELD_HALF_Z - ArenaBoundary.BASE_RADIUS


func _random_direction() -> Vector3:
	var direction := Vector3(randf_range(-1, 1), randf_range(-1, 1), randf_range(-1, 1))
	return direction.normalized() if direction.length_squared() > 0.001 else Vector3.FORWARD


func _place_body(body: RigidBody3D, to: Transform3D, linear_velocity: Vector3, angular_velocity: Vector3) -> void:
	# Deferred: a RigidBody3D transform can't be set mid-physics-step
	body.set_deferred("global_transform", to)
	body.set_deferred("linear_velocity", linear_velocity)
	body.set_deferred("angular_velocity", angular_velocity)