CosmicClash/Game/scripts/ship_ai_controller.gd

class_name ShipAIController
extends AIController3D

# Training-side bridge between godot_rl_agents and a ship. This is the only
# class that touches plugin types (AIController3D / the Sync node protocol) —
# everything else stays behind the ShipController seam: actions received from
# the trainer are written into an RLShipController, which the ship pulls like
# any other controller.
#
# Action space is ShipAction verbatim: 6 continuous axes (thrust xyz,
# rotation xyz, each -1..1) + binary turbo. ShipAction axes are ship-local
# (body frame), so they need no team mirroring — only observations do
# (see ShipObservations.canon).

# Reward shaping weights. Dense terms accrue per physics tick (60 sim-ticks
# per sim-second); event terms fire once. Exported so tuning needs no code
# edits. Goal rewards are added by TrainingMode, which owns goal events.
@export var ball_touch_reward := 0.4
# Ball touches pay out at most once per this many physics ticks (1 sim-
# second at 60). Run07 lesson: body_entered re-fires on every micro-
# separation, so pinning the ball against a surface farmed ~2 touches/s —
# outearning every other term while the goal rate fell. The cooldown keeps
# touches a stepping-stone signal instead of the objective. Halved again
# after run01-vs-run02 eval (training/eval_history.json) came back 87.5%
# draws: even at 1 touch/s, a full episode's worth of touches could still
# outweigh TrainingMode's goal_reward, so scoring and ending the episode
# early was never worth it. See goal_reward's comment for the other half of
# this fix.
@export var ball_touch_cooldown_ticks := 60
# A touch pays out scaled by how goal-directed it was — full ball_touch_reward
# when the post-touch ball velocity points straight at the attack goal, down
# to this floor when it doesn't (0 = only goal-directed touches pay at all).
# Without this, any contact paid the same regardless of direction, so batting
# the ball anywhere counted the same as an actual shot on goal — reinforcing
# possession, not scoring. The floor keeps a purely defensive touch (e.g.
# clearing a shot away from your own goal) worth something as a stepping
# stone, matching ball_touch_cooldown_ticks's existing "stepping-stone, not
# the objective" framing.
@export_range(0.0, 1.0) var ball_touch_direction_floor := 0.3
@export var velocity_to_ball_weight := 0.02
@export var ball_velocity_to_goal_weight := 0.004
# Per-tick penalty scaled by distance to the ball (full value at the arena's
# far diagonal, 0 on top of the ball). Run04 lesson: with idling worth a flat
# 0, camping in a corner strictly dominated risking the wall/tilt penalties
# to chase the ball — this makes "do nothing far from the ball" the worst
# option instead of the safest. A penalty, not a proximity bonus, so orbiting
# the ball farms nothing.
@export var ball_distance_penalty := 0.002
# Per-tick penalty while pressed against a side wall, end wall, or the
# ceiling — NOT the floor (run03 lesson: taxing floor contact punishes the
# ship's natural low flight and drowns every other signal). At 60 ticks per
# sim-second this is -0.15/s. Halved for run05: the ball lives near walls,
# and the old -0.3/s made the productive region of the pitch aversive
# relative to the (then far weaker) ball-seeking shaping.
@export var wall_contact_penalty := 0.0025
# Per-tick penalty for not being upright, scaled by tilt: 0 when flat, full
# value (-0.12/s) when inverted. A penalty rather than an upright bonus so a
# flat, idle ship farms nothing.
@export var tilt_penalty := 0.002
# Per-tick bonus for own speed: 0 stationary, full value (+0.24/s) at
# max_speed. Run07 lesson: after the kickoff flurry both ships parked next to
# a cornered ball — with every other dense term near zero there, standing
# still was a rest state. Sized well below velocity_to_ball_weight so flying
# fast toward the ball still beats flying fast anywhere else.
@export var speed_reward_weight := 0.004
# Flat per-tick cost (-0.06/s, -1.8 over a full 30s episode) applied
# regardless of position or behaviour. Every other dense term can be farmed
# indefinitely by an episode that never ends in a goal; this one can't — it
# only stops accruing once the episode does, via a goal or the timeout. That
# makes running the clock out strictly worse than scoring as soon as a
# chance appears, instead of a free way to keep collecting dense reward.
@export var time_penalty := 0.001
# Per-tick penalty scaled by height above the floor (0 on the floor, full
# value at the arena's ceiling) — distinct from the locomotion mask, which
# only discards *thrust*-driven vertical/pitch-roll input; a masked ship can
# still be launched airborne by collisions (ball impacts, ship-vs-ship
# knockback, the wall/ceiling surface-pull field), and nothing previously
# penalized time spent up there. Default 0 (off) so ordinary runs are
# unaffected; the floor-lock curriculum stage turns it on.
@export var airborne_penalty := 0.0

# Locomotion curriculum: scales how much of the corresponding action axes
# actually reaches the ship, from 0.0 (fully discarded, grounded-only) to
# 1.0 (full effect) — this scales the *effect* of thrust.y/rotation.x/
# rotation.z in set_action, not the action space's shape: the policy always
# outputs values for these axes (always contributing to PPO's entropy/log-
# prob), they're just attenuated here, so checkpoints stay resumable across
# ramp values.
#
# A hard 0/1 flip (the original bool mask) let PPO's action-distribution
# std collapse to ~0.13-0.15 within the first ~10% of steps, before the
# policy ever meaningfully explored the newly-unmasked axes — 3 independent
# 240M-step attempts at the all-or-nothing flip all landed at a stable
# ~28-32% win rate vs curric-s5-aggression (see TRAINING.md's generation 3
# section). A gradual ramp across several short curriculum stages, each
# resuming from the previous ramp value's checkpoint, lets the policy adopt
# each axis incrementally instead of all at once.
@export_range(0.0, 1.0) var vertical_ramp := 1.0
@export_range(0.0, 1.0) var pitch_roll_ramp := 1.0

# Contact normals with y above this are floor contact (exempt from the wall
# penalty); below it they read as wall (sideways) or ceiling (downward).
const FLOOR_NORMAL_MIN_Y := 0.7

# Longest possible ship-to-ball separation: the enclosure's interior diagonal.
# Normalizes ball_distance_penalty so its export is the worst-case per-tick cost.
const MAX_BALL_DISTANCE := sqrt(
	(2.0 * ArenaBoundary.INNER_HALF_X) ** 2
	+ (2.0 * ArenaBoundary.INNER_HALF_Z) ** 2
	+ ArenaBoundary.INNER_HEIGHT ** 2
)

var ship: Ship
var rl_controller: RLShipController
var ball: RigidBody3D
var opponent: Ship
var attack_goal_position: Vector3

# Set directly by TrainingMode (_on_goal_scored / the timeout branch in
# _physics_process) at the same time as `done = true`. Deliberately NOT
# cleared in reset(): TrainingMode's _reset_episode() (which calls reset())
# runs synchronously, immediately after done is set, before the Sync node
# ever reads get_info()/get_done() for that terminal tick — clearing it here
# would wipe the value that read needs. Both call sites always overwrite
# (true on goal, false on timeout) rather than toggle, so no reset is needed.
var goal_scored_this_episode := false

var _ticks_since_ball_touch := 1 << 30  # large so the first touch always pays


# Wire up references after the ship is spawned. `attack_goal` is the goal
# this ship scores into (goal.team == opponent's team).
func setup(p_ship: Ship, p_rl_controller: RLShipController, p_ball: RigidBody3D, p_opponent: Ship, p_attack_goal_position: Vector3) -> void:
	ship = p_ship
	rl_controller = p_rl_controller
	ball = p_ball
	opponent = p_opponent
	attack_goal_position = p_attack_goal_position
	init(ship)

	# Contact monitoring for the ball-touch reward (training-only cost;
	# the shipped game leaves contact_monitor off).
	ship.contact_monitor = true
	ship.max_contacts_reported = 8
	ship.body_entered.connect(_on_ship_body_entered)


func get_obs() -> Dictionary:
	return {"obs": ShipObservations.build(ship, opponent, ball, attack_goal_position)}


func get_reward() -> float:
	return reward


# Symmetric across both self-play agents: reports whether this episode ended
# in a goal at all, not which team scored — a clean "goal rate" signal
# distinct from rollout/ep_rew_mean, which mixes this with dense shaping
# (ball chasing/touching). See train.py's GoalRateCallback.
func get_info() -> Dictionary:
	return {"goal_scored": goal_scored_this_episode}


func get_action_space() -> Dictionary:
	return {
		"thrust": {"size": 3, "action_type": "continuous"},
		"rotation": {"size": 3, "action_type": "continuous"},
		"turbo": {"size": 2, "action_type": "discrete"},
	}


func set_action(action) -> void:
	var thrust: Array = action["thrust"]
	var rot: Array = action["rotation"]
	var thrust_y: float = thrust[1] * vertical_ramp
	var pitch: float = rot[0] * pitch_roll_ramp
	var roll: float = rot[2] * pitch_roll_ramp
	rl_controller.action.thrust = Vector3(thrust[0], thrust_y, thrust[2])
	rl_controller.action.rotation = Vector3(pitch, rot[1], roll)
	rl_controller.action.turbo = int(action["turbo"]) == 1


func reset():
	super()
	_ticks_since_ball_touch = 1 << 30


func _physics_process(delta):
	super(delta)
	if not is_instance_valid(ship) or not is_instance_valid(ball):
		return
	_ticks_since_ball_touch += 1

	# Flat time cost — see time_penalty.
	reward -= time_penalty

	# Dense shaping: own velocity toward the ball
	var to_ball := ball.global_position - ship.global_position
	if to_ball.length_squared() > 0.0001:
		var closing_speed := ship.linear_velocity.dot(to_ball.normalized())
		reward += velocity_to_ball_weight * closing_speed / ship.max_speed

	# Dense penalty: distance to the ball, so idling far away bleeds reward
	# instead of scoring a safe zero (see ball_distance_penalty).
	if ball_distance_penalty > 0.0:
		reward -= ball_distance_penalty * to_ball.length() / MAX_BALL_DISTANCE

	# Dense bonus: own speed, so hovering in place is never a rest state
	# (see speed_reward_weight).
	if speed_reward_weight > 0.0:
		reward += speed_reward_weight * ship.linear_velocity.length() / ship.max_speed

	# Dense shaping: ball velocity toward the goal we attack
	var ball_to_goal := attack_goal_position - ball.global_position
	if ball_to_goal.length_squared() > 0.0001:
		var ball_progress := ball.linear_velocity.dot(ball_to_goal.normalized())
		reward += ball_velocity_to_goal_weight * ball_progress / ShipObservations.BALL_SPEED_SCALE

	# Dense penalty: every tick spent pressed against a wall or the ceiling
	# (contact monitoring is already on for the ball-touch reward). Ships
	# bumping each other, the ball, or the floor is fine. The boundary is one
	# body, so the contact normal tells us which surface: floor contact
	# pushes the ship up (+Y), walls push sideways, the ceiling down.
	if wall_contact_penalty > 0.0 and _wall_or_ceiling_contact():
		reward -= wall_contact_penalty

	# Dense penalty: tilt away from upright (0 flat, max when inverted) —
	# discourages ending up on a side or roof without rewarding idleness.
	if tilt_penalty > 0.0:
		var uprightness: float = ship.global_transform.basis.y.dot(Vector3.UP)
		reward -= tilt_penalty * (1.0 - uprightness) * 0.5

	# Dense penalty: height above the floor (see airborne_penalty). The
	# floor sits at world y = 0 (see training_mode.gd's FIELD_MIN_Y/
	# _escaped bounds); normalized so the worst case is pinned at the
	# ceiling.
	if airborne_penalty > 0.0:
		var height := maxf(ship.global_position.y, 0.0)
		reward -= airborne_penalty * height / ArenaBoundary.INNER_HEIGHT


func _wall_or_ceiling_contact() -> bool:
	var state := PhysicsServer3D.body_get_direct_state(ship.get_rid())
	if state == null:
		return false
	for i in state.get_contact_count():
		if not state.get_contact_collider_object(i) is ArenaBoundary:
			continue
		# Normal points from the surface into the ship: floor ≈ +Y (exempt),
		# anything flatter or downward is a wall or the ceiling.
		if state.get_contact_local_normal(i).y < FLOOR_NORMAL_MIN_Y:
			return true
	return false


func _on_ship_body_entered(body: Node) -> void:
	if not body.is_in_group("ball") or _ticks_since_ball_touch < ball_touch_cooldown_ticks:
		return
	# Contact-signal ordering means ball.linear_velocity here already reflects
	# the collision impulse from this touch, not the pre-touch velocity.
	var alignment := 0.0
	var to_goal := attack_goal_position - ball.global_position
	if to_goal.length_squared() > 0.0001 and ball.linear_velocity.length_squared() > 0.0001:
		alignment = clampf(ball.linear_velocity.normalized().dot(to_goal.normalized()), 0.0, 1.0)
	reward += ball_touch_reward * lerpf(ball_touch_direction_floor, 1.0, alignment)
	_ticks_since_ball_touch = 0