Apollo3D_SE/agent/agent.py

from dataclasses import Field
import logging
from typing import Mapping

import numpy as np
from utils.math_ops import MathOps
from world.commons.field import FIFAField, HLAdultField, Soccer7vs7Field
from world.commons.play_mode import PlayModeEnum, PlayModeGroupEnum


logger = logging.getLogger()


class Agent:
    """
    Responsible for deciding what the agent should do at each moment.

    This class is called every simulation step to update the agent's behavior
    based on the current state of the world and game conditions.
    """

    BEAM_POSES: Mapping[type[Field], Mapping[int, tuple[float, float, float]]] ={
        FIFAField: {
            1: (2.1, 0, 0),
            2: (22.0, 12.0, 0),
            3: (22.0, 4.0, 0),
            4: (22.0, -4.0, 0),
            5: (22.0, -12.0, 0),
            6: (15.0, 0.0, 0),
            7: (4.0, 16.0, 0),
            8: (11.0, 6.0, 0),
            9: (11.0, -6.0, 0),
            10: (4.0, -16.0, 0),
            11: (7.0, 0.0, 0),
        },
        HLAdultField: {
            1: (7.0, 0.0, 0),
            2: (2.0, -1.5, 0),
            3: (2.0, 1.5, 0),
        },
        Soccer7vs7Field: {
            1: (2.1, 0, 0),
            2: (22.0, 12.0, 0),
            3: (22.0, 4.0, 0),
            4: (22.0, -4.0, 0),
            5: (22.0, -12.0, 0),
            6: (15.0, 0.0, 0),
            7: (4.0, 16.0, 0)
        }
    }

    TEAMMATE_INFO_MAX_AGE: float = 0.8
    SUPPORT_DISTANCE_FROM_BALL: float = 1.2

    def __init__(self, agent):
        """
        Creates a new DecisionMaker linked to the given agent.

        Args:
            agent: The main agent that owns this DecisionMaker.
        """
        from agent.base_agent import Base_Agent  # type hinting

        self.agent: Base_Agent = agent
        self.is_getting_up: bool = False
        self.is_kicking: bool = False
        self.current_shot_target_point: np.ndarray | None = None
        self.target_point: np.ndarray = np.array([0.0, 0.0])
        self.my_pos: np.ndarray = self.agent.world.global_position[:2]

    def update_current_behavior(self) -> None:
        """
        Chooses what the agent should do in the current step.

        This function checks the game state and decides which behavior
        or skill should be executed next.
        """

        if self.agent.world.playmode is PlayModeEnum.GAME_OVER:
            return

        if self.agent.world.playmode_group in (
            PlayModeGroupEnum.ACTIVE_BEAM,
            PlayModeGroupEnum.PASSIVE_BEAM,
        ):
            self.agent.server.commit_beam(
                pos2d=self.BEAM_POSES[type(self.agent.world.field)][self.agent.world.number][:2],
                rotation=self.BEAM_POSES[type(self.agent.world.field)][self.agent.world.number][2],
            )

        if self.is_getting_up or self.agent.skills_manager.is_ready(skill_name="GetUp"):
            self.is_getting_up = not self.agent.skills_manager.execute(skill_name="GetUp")

        elif self.agent.world.playmode is PlayModeEnum.PLAY_ON:
            active_player_number = self.get_active_player_number()
            if active_player_number == self.agent.world.number:
                # self.carry_ball()
                self.target_point = self.my_pos + np.array([4.0, 0.0])
                self.kick_ball()
            else:
                # self.execute_support_behavior()
                self.agent.skills_manager.execute("Neutral")
        elif self.agent.world.playmode in (PlayModeEnum.BEFORE_KICK_OFF, PlayModeEnum.THEIR_GOAL, PlayModeEnum.OUR_GOAL):
            self.agent.skills_manager.execute("Neutral")
        else:
            # self.execute_support_behavior()
            self.kick_ball()

        self.agent.robot.commit_motor_targets_pd()

    def carry_ball(self):
        """
        Basic example of a behavior: moves the robot toward the goal while handling the ball.
        """
        target_point = self.target_point
        ball_pos = self.agent.world.ball_pos[:2]
        my_pos = self.agent.world.global_position[:2]

        if self.kick_ball(target_point=target_point):
            return

        ball_to_target = target_point - ball_pos
        bt_norm = np.linalg.norm(ball_to_target)
        if bt_norm == 0:
            return
        ball_to_target_dir = ball_to_target / bt_norm

        dist_from_ball_to_start_carrying = 0.30
        carry_ball_pos = ball_pos - ball_to_target_dir * dist_from_ball_to_start_carrying

        my_to_ball = ball_pos - my_pos
        my_to_ball_norm = np.linalg.norm(my_to_ball)
        if my_to_ball_norm == 0:
            my_to_ball_dir = np.zeros(2)
        else:
            my_to_ball_dir = my_to_ball / my_to_ball_norm

        cosang = np.dot(my_to_ball_dir, ball_to_target_dir)
        cosang = np.clip(cosang, -1.0, 1.0)
        angle_diff = np.arccos(cosang)

        ANGLE_TOL = np.deg2rad(7.5)
        aligned = (my_to_ball_norm > 1e-6) and (angle_diff <= ANGLE_TOL)

        behind_ball = np.dot(my_pos - ball_pos, ball_to_target_dir) < 0
        desired_orientation = MathOps.vector_angle(ball_to_target)

        if not aligned or not behind_ball:
            self.agent.skills_manager.execute(
                "Walk",
                target_2d=carry_ball_pos,
                is_target_absolute=True,
                orientation=None if np.linalg.norm(my_pos - carry_ball_pos) > 2 else desired_orientation
            )
        else:
            self.agent.skills_manager.execute(
                "Walk",
                target_2d=target_point,
                is_target_absolute=True,
                orientation=desired_orientation
            )

    def kick_ball(self, target_point=None):
        if target_point is None:
            target_point = self.target_point

        target_point = np.asarray(target_point, dtype=float)[:2]
        ball_pos = self.agent.world.ball_pos[:2]
        my_pos = self.agent.world.global_position[:2]

        ball_to_target = target_point - ball_pos
        ball_to_target_dist = np.linalg.norm(ball_to_target)
        if ball_to_target_dist <= 1e-6:
            return False
        if self.is_kicking:
            if self.current_shot_target_point is None:
                self.current_shot_target_point = target_point

            active_target = self.current_shot_target_point
            active_distance = np.linalg.norm(active_target - ball_pos)
            self.is_kicking = not self.agent.skills_manager.execute(
                "Kick", distance=active_distance
            )
            if not self.is_kicking:
                self.current_shot_target_point = None
            return True

        my_to_ball = ball_pos - my_pos
        my_to_ball_dist = np.linalg.norm(my_to_ball)

        ball_to_target_dir = ball_to_target / ball_to_target_dist
        behind_ball = np.dot(my_pos - ball_pos, ball_to_target_dir) < 0
        me_to_ball_dir = ball_pos - my_pos
        me_to_ball_dir_norm = np.linalg.norm(me_to_ball_dir)

        desired_orientation = MathOps.vector_angle(me_to_ball_dir)
        current_orientation = self.agent.robot.global_orientation_euler[2]
        orientation_error = abs(
            MathOps.normalize_deg(desired_orientation - current_orientation)
        )

        lateral_error_to_shot_line, _ = MathOps.distance_point_to_line(
            p=my_pos,
            a=ball_pos,
            b=target_point,
        )

        kick_offset = MathOps.rotate_2d_vec(np.array([0.0, -0.1]), desired_orientation)

        close_to_ball = my_to_ball_dist <= 0.24
        good_orientation = orientation_error <= 8.0
        good_lateral_offset = lateral_error_to_shot_line <= 0.16

        can_shoot = (
            close_to_ball
            and behind_ball
            and good_orientation
            and good_lateral_offset
        )

        if not can_shoot:
            prepare_offset = 0.1
            prepare_pos = ball_pos - ball_to_target_dir * prepare_offset + kick_offset

            self.agent.skills_manager.execute(
                "Walk",
                target_2d=prepare_pos,
                is_target_absolute=True,
                orientation=desired_orientation,
            )
            return True

        self.current_shot_target_point = target_point
        self.is_kicking = not self.agent.skills_manager.execute(
            "Kick", distance=ball_to_target_dist
        )
        if not self.is_kicking:
            self.current_shot_target_point = None
        return True

    def get_active_player_number(self) -> int:
        world = self.agent.world
        ball_pos = world.ball_pos[:2]
        server_time = world.server_time

        my_num = world.number
        my_pos = world.global_position[:2]
        best_player_number = my_num
        best_distance = np.linalg.norm(my_pos - ball_pos)

        for teammate_number, teammate in enumerate(world.our_team_players, start=1):
            if teammate_number == my_num:
                continue

            if teammate.last_seen_time is None or server_time is None:
                continue

            info_age = server_time - teammate.last_seen_time
            if info_age > self.TEAMMATE_INFO_MAX_AGE:
                continue

            teammate_pos = teammate.position[:2]
            teammate_distance = np.linalg.norm(teammate_pos - ball_pos)

            if teammate_distance + 1e-6 < best_distance:
                best_distance = teammate_distance
                best_player_number = teammate_number
            elif abs(teammate_distance - best_distance) <= 1e-6:
                best_player_number = min(best_player_number, teammate_number)

        return best_player_number

    def execute_support_behavior(self) -> None:
        world = self.agent.world
        my_pos = world.global_position[:2]
        ball_pos = world.ball_pos[:2]
        target_point = self.target_point

        ball_to_target = target_point - ball_pos
        bt_norm = np.linalg.norm(ball_to_target)
        if bt_norm <= 1e-6:
            ball_to_target_dir = np.array([1.0, 0.0])
        else:
            ball_to_target_dir = ball_to_target / bt_norm

        support_pos = ball_pos - ball_to_target_dir * self.SUPPORT_DISTANCE_FROM_BALL
        support_orientation = MathOps.vector_angle(ball_pos - my_pos)

        self.agent.skills_manager.execute(
            "Walk",
            target_2d=support_pos,
            is_target_absolute=True,
            orientation=support_orientation,
        )