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)
        }
    }

    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

    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:
            self.carry_ball()
        elif self.agent.world.playmode in (PlayModeEnum.BEFORE_KICK_OFF, PlayModeEnum.THEIR_GOAL, PlayModeEnum.OUR_GOAL):
            self.agent.skills_manager.execute("Neutral")
        else:
            self.carry_ball()

        self.agent.robot.commit_motor_targets_pd()

    def carry_ball(self):
        """
        Minimal catch-ball behavior.

        All players share the same logic:
        1. approach a point behind the ball
        2. reposition with a lateral offset if they are close but not yet behind it
        3. push the ball forward once they are aligned
        """
        their_goal_pos = self.agent.world.field.get_their_goal_position()[:2]
        ball_pos = self.agent.world.ball_pos[:2]
        my_pos = self.agent.world.global_position[:2]

        ball_to_goal = their_goal_pos - ball_pos
        bg_norm = np.linalg.norm(ball_to_goal)
        if bg_norm <= 1e-6:
            ball_to_goal_dir = np.array([1.0, 0.0])
        else:
            ball_to_goal_dir = ball_to_goal / bg_norm

        lateral_dir = np.array([-ball_to_goal_dir[1], ball_to_goal_dir[0]])

        back_offset = 0.40
        side_offset = 0.35
        push_distance = 0.80
        approach_distance = 0.90
        push_start_distance = 0.55
        behind_margin = 0.08
        angle_tolerance = np.deg2rad(20.0)

        behind_point = ball_pos - ball_to_goal_dir * back_offset
        push_target = ball_pos + ball_to_goal_dir * push_distance

        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_goal_dir)
        cosang = np.clip(cosang, -1.0, 1.0)
        angle_diff = np.arccos(cosang)
        aligned = (my_to_ball_norm > 1e-6) and (angle_diff <= angle_tolerance)

        behind_ball = np.dot(my_pos - ball_pos, ball_to_goal_dir) < -behind_margin
        desired_orientation = MathOps.vector_angle(ball_to_goal)

        lateral_sign = np.sign(np.cross(ball_to_goal_dir, my_to_ball_dir))
        if lateral_sign == 0:
            lateral_sign = 1.0 if (my_pos[1] - ball_pos[1]) >= 0 else -1.0

        reposition_point = behind_point + lateral_dir * lateral_sign * side_offset

        if my_to_ball_norm > approach_distance:
            target_2d = behind_point
            orientation = None
        elif not behind_ball:
            target_2d = reposition_point
            orientation = None if np.linalg.norm(my_pos - reposition_point) > 0.8 else desired_orientation
        elif not aligned and my_to_ball_norm > push_start_distance:
            target_2d = behind_point
            orientation = desired_orientation
        else:
            target_2d = push_target
            orientation = desired_orientation

        if np.linalg.norm(target_2d - my_pos) <= 1e-4:
            target_2d = my_pos + ball_to_goal_dir * 0.30

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