Aimbot-ParallelEnv/Assets/RewardFunction.cs

using System;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;

public class RewardFunction : MonoBehaviour
{
    public enum EndType
    { Win, Lose, Running, Num };

    [SerializeField] private GameObject parameterContainerObj;
    [SerializeField] private GameObject sceneBlockContainerObj;
    [SerializeField] private GameObject targetControllerObj;
    [SerializeField] private GameObject environmentUIObj;
    [SerializeField] private GameObject commonParameterContainerObj;
    [SerializeField] private GameObject enemyContainerObj;

    private GameObject agentObj;
    private Camera fpsCam;
    private CommonParameterContainer commonParamCon;
    private SceneBlockContainer sceneBlockCon;
    private ParameterContainer paramCon;
    private TargetController targetCon;
    private EnvironmentUIControl envUICon;
    private AgentController agentCon;
    private RaySensors raySensors;

    private bool firstRewardFlag = false;
    private float lastDistance;
    private float lastEnemyFacingDistance = 0f; // record last enemy facing minimum distance
    private float lastTargetFacingDistance = 0f; // record last target facing minimum distance
    private List<float> spinRecord = new List<float>();

    private void Start()
    {
        agentObj = gameObject;
        agentCon = agentObj.GetComponent<AgentController>();
        fpsCam = agentCon.fpsCam;
        commonParamCon = commonParameterContainerObj.GetComponent<CommonParameterContainer>();
        paramCon = parameterContainerObj.GetComponent<ParameterContainer>();
        sceneBlockCon = sceneBlockContainerObj.GetComponent<SceneBlockContainer>();
        targetCon = targetControllerObj.GetComponent<TargetController>();
        envUICon = environmentUIObj.GetComponent<EnvironmentUIControl>();
        raySensors = GetComponent<RaySensors>();
    }

    /// <summary>
    /// Calculates the reward value.
    /// </summary>
    /// <param name="sceneReward">Reward value from the scene.</param>
    /// <param name="mouseX">Movement amount of the mouse along the X-axis.</param>
    /// <param name="movement">Movement of discrete.</param>
    /// <param name="shootState">State of the shooting action.</param>
    /// <returns>Returns the calculated total reward value.</returns>
    /// <remarks>
    /// This method calculates the total reward based on the provided parameters,
    /// taking into account rewards for enemy kills, shooting actions, facing reward,
    /// and penalties such as spin and movement.
    /// </remarks>
    public float RewardCalculate(float sceneReward, float mouseX, float movement, int shootState)
    {
        float epreward = 0f;
        // Got kill point reward
        if (agentCon.enemyKillCount > 0)
        {
            for (int i = 0; i < agentCon.enemyKillCount; i++)
            {
                // get
                epreward += KillReward(agentCon.killEnemyPosition);
            }
            agentCon.enemyKillCount = 0;
        }
        else
        {
            agentCon.enemyKillCount = 0;
        }
        // Shoot action reward
        epreward += Ballistic(shootState) + sceneReward;
        // facing reward
        epreward += FacingReward();
        // Penalty
        // spin penalty
        spinRecord.Add(mouseX);
        if (spinRecord.Count >= commonParamCon.spinRecordMax)
        {
            spinRecord.RemoveAt(0);
        }
        float spinPenaltyReward = Math.Abs(spinRecord.ToArray().Sum() * commonParamCon.spinPenalty);
        if (spinPenaltyReward >= commonParamCon.spinPenaltyThreshold)
        {
            epreward -= spinPenaltyReward;
        }
        else
        {
            epreward -= Math.Abs(mouseX) * commonParamCon.mousePenalty;
        }
        // move penalty
        if (movement != 0)
        {
            epreward -= commonParamCon.movePenalty;
        }
        return epreward;
    }

    /// <summary>
    /// Calculates the reward value for shooting actions.
    /// </summary>
    /// <param name="shootState">State value of the shooting action.</param>
    /// <returns>Returns the reward value associated with shooting.</returns>
    /// <remarks>
    /// This method calculates the reward value based on the shooting state and other related conditions,
    /// such as whether the enemy was hit, whether the shot was towards the target area, and whether the gun was ready to shoot.
    /// </remarks>
    private float Ballistic(int shootState)
    {
        Vector3 point = new Vector3(fpsCam.pixelWidth / 2, fpsCam.pixelHeight / 2, 0);// start position
        Ray ray = fpsCam.ScreenPointToRay(point);
        RaycastHit hit;
        // Debug.DrawRay(centerRay.origin, centerRay.direction * 100, Color.blue);
        // Mouse Pressed
        if (shootState != 0 && agentCon.gunReadyToggle == true)
        {
            agentCon.lastShootTime = Time.time;
            if (Physics.Raycast(ray, out hit, 100))
            {
                if (hit.collider.tag != agentCon.myTag && hit.collider.tag != "Wall" && hit.collider.tag != "Untagged")
                {
                    // kill enemy
                    GameObject gotHitObj = hit.transform.gameObject;
                    gotHitObj.GetComponent<States>().ReactToHit(commonParamCon.damage, gameObject);
                    shootState = 0;
                    return HitEnemyReward(gotHitObj.transform.position);
                }
            }
            if (targetCon.targetType == Targets.Attack)
            {
                // while if attack mode
                float targetDis = Vector3.Distance(sceneBlockCon.nowBlock.transform.position, transform.position);
                if (targetDis <= raySensors.viewDistance)
                {
                    // Debug.DrawRay(new Vector3(0,0,0), viewPoint, Color.red);
                    if (Vector3.Distance(ray.origin + (ray.direction * targetDis), sceneBlockCon.nowBlock.transform.position) <= sceneBlockCon.nowBlock.firebasesAreaDiameter / 2)
                    {
                        // im shooting at target but didn't hit enemy
                        // Debug.DrawRay(centerRay.origin, viewPoint-centerRay.origin, Color.blue);
                        return commonParamCon.shootTargetAreaReward;
                    }
                }
            }
            shootState = 0;
            return commonParamCon.shootReward;
        }
        else if (shootState != 0 && agentCon.gunReadyToggle == false)
        {
            // shoot without ready
            shootState = 0;
            return commonParamCon.shootWithoutReadyReward;
        }
        else
        {
            // do not shoot
            shootState = 0;
            return commonParamCon.nonReward;
        }
    }

    /// <summary>
    /// Retrieves the reward value based on the character's facing direction.
    /// </summary>
    /// <returns>Returns the reward value for the facing direction.</returns>
    /// <remarks>
    /// This method calculates a reward value based on the relationship between the character's facing direction and the target.
    /// in free mode, if the character is facing an enemy, the reward is a fixed value
    /// in attack mode, the reward depends on the distance between the character and the target, among other factors.
    /// </remarks>
    private float FacingReward()
    {
        Vector3 screenCenter = new Vector3(fpsCam.pixelWidth / 2, fpsCam.pixelHeight / 2, 0);
        Vector3 screenLeft = new Vector3(0, fpsCam.pixelHeight / 2, 0);

        Ray centerRay = fpsCam.ScreenPointToRay(screenCenter);
        Ray leftRay = fpsCam.ScreenPointToRay(screenLeft);

        switch (targetCon.targetType)
        {
            case Targets.Free:
                return FacingRewardFree(centerRay);

            case Targets.Attack:
                return FacingRewardAttack(centerRay, leftRay);

            case Targets.Go:
                return FacingRewardGo(centerRay, leftRay);

            case Targets.Stay:
                // stay mode has no facing reward
                return 0f;

            default:
                Debug.LogError("Wrong target type");
                return 0f;
        }
    }

    private float FacingRewardFree(Ray centerRay)
    {
        float nowReward = 0;
        float enemyFacingDistance = 0f;
        bool isFacingtoEnemy = false;
        RaycastHit hit;
        if (Physics.Raycast(centerRay, out hit, 100))
        {
            // facing to an enemy
            if (hit.collider.tag != agentCon.myTag && hit.collider.tag != "Wall")
            {
                nowReward = commonParamCon.facingReward;
                isFacingtoEnemy = true;
            }
        }
        if (raySensors.inViewEnemies.Count > 0 && !isFacingtoEnemy)
        {
            // have enemy in view
            List<float> projectionDis = new List<float>();
            foreach (GameObject theEnemy in raySensors.inViewEnemies)
            {
                // for each enemy in view
                Vector3 projection = Vector3.Project(theEnemy.transform.position - transform.position, (centerRay.direction * 10));
                Vector3 verticalToRay = transform.position + projection - theEnemy.transform.position;
                projectionDis.Add(verticalToRay.magnitude);
                // Debug.Log("enemy!" + verticalToRay.magnitude);
                // Debug.DrawRay(transform.position, (centerRay.direction * 100), Color.cyan);
                // Debug.DrawRay(transform.position, theEnemy.transform.position - transform.position, Color.yellow);
                // Debug.DrawRay(transform.position, projection, Color.blue);
                // Debug.DrawRay(theEnemy.transform.position, verticalToRay, Color.magenta);
            }
            enemyFacingDistance = projectionDis.Min();
            if (enemyFacingDistance <= lastEnemyFacingDistance)
            {
                // closing to enemy
                nowReward = 1 / MathF.Sqrt(commonParamCon.facingInviewEnemyDisCOEF * enemyFacingDistance + 0.00001f);
            }
            else
            {
                nowReward = 0;
            }
            // enemy in view Reward
            lastEnemyFacingDistance = enemyFacingDistance;
            if (nowReward >= commonParamCon.facingReward) nowReward = commonParamCon.facingReward; // limit
            if (nowReward <= -commonParamCon.facingReward) nowReward = -commonParamCon.facingReward; // limit
            // Debug.Log("ninimum = " + nowReward);
        }
        return nowReward;
    }

    private float FacingRewardGo(Ray centerRay, Ray leftRay)
    {
        float nowReward = 0;
        float camCenterToFireBase;
        float camCenterToViewEdge;
        (camCenterToFireBase, camCenterToViewEdge, _) = CameraCenterToFireBaseAndViewEdge(centerRay, leftRay);

        // goto mode
        if (camCenterToFireBase <= camCenterToViewEdge)
        {
            // fireArea is in view
            nowReward = commonParamCon.facingReward;
        }
        else
        {
            nowReward = 0;
        }
        return nowReward;
    }

    private float FacingRewardAttack(Ray centerRay, Ray leftRay)
    {
        float nowReward = 0;
        float camCenterToFireBase;
        float targetDis;
        (camCenterToFireBase, _, targetDis) = CameraCenterToFireBaseAndViewEdge(centerRay, leftRay);
        // attack mode
        if (targetDis <= raySensors.viewDistance)
        {
            // Debug.DrawRay(new Vector3(0,0,0), viewPoint, Color.red);
            // while center of screen between target's distance is lower than firebasesAreaDiameter
            // while facing to target
            if (camCenterToFireBase <= sceneBlockCon.nowBlock.firebasesAreaDiameter / 2)
            {
                // Debug.DrawRay(centerRay.origin, viewPoint-centerRay.origin, Color.blue);
                nowReward = commonParamCon.facingReward;
            }
            else
            {
                // while not facing to target
                nowReward = (lastTargetFacingDistance - camCenterToFireBase) * commonParamCon.facingTargetReward;
            }
        }
        // update lastTargetFacingDistance
        lastTargetFacingDistance = camCenterToFireBase;
        return nowReward;
    }

    private (float, float, float) CameraCenterToFireBaseAndViewEdge(Ray centerRay, Ray leftRay)
    {
        // target fireBaseArea Position, turen y to camera's y
        Vector3 fireBaseArea = sceneBlockCon.nowBlock.fireBasesAreaObj.transform.position;
        fireBaseArea.y = fpsCam.transform.position.y;

        // my position, turn y to camera's y
        // Debug.DrawRay(fpsCam.transform.position, centerRay.direction * 100, Color.blue);
        Vector3 myposition = transform.position;
        myposition.y = fpsCam.transform.position.y;

        // Target to Agent distance
        //Debug.DrawLine(fireBaseArea, myposition, Color.red);
        float targetDis = Vector3.Distance(fireBaseArea, myposition);

        // point in centerRay and leftRay which distance is targetDis from camera center
        Vector3 pointInCenterRay = fpsCam.transform.position + (centerRay.direction * targetDis);
        Vector3 pointInLeftRay = fpsCam.transform.position + (leftRay.direction * targetDis);

        // center of screen to target's distance
        // Debug.DrawLine(pointInCenterRay, fireBaseArea,Color.green);
        float camCenterToFireBase = Vector3.Distance(pointInCenterRay, fireBaseArea);

        // left of screen to target's distance
        // Debug.DrawLine(pointInLeftRay, pointInCenterRay, Color.yellow);
        float camCenterToViewEdge = Vector3.Distance(pointInLeftRay, pointInCenterRay);
        return (camCenterToFireBase, camCenterToViewEdge, targetDis);
    }

    /// <summary>
    /// Checks the game's end state and retrieves rewards.
    /// </summary>
    /// <returns>A tuple containing the game's end type, current reward, and final reward.
    /// 1 = success,2 = overtime,0 = notover</returns>
    public (int, float, float) CheckOverAndRewards()
    {
        int endTypeInt = 0;
        float nowReward = 0;
        float endReward = 0;
        switch (targetCon.targetType)
        {
            case Targets.Go:
                // goto
                (endTypeInt, nowReward, endReward) = CheckOverAndRewardsGo();
                break;

            case Targets.Attack:
                // attack
                (endTypeInt, nowReward, endReward) = CheckOverAndRewardsAttack();
                break;

            case Targets.Defence:
                //defence
                (endTypeInt, nowReward, endReward) = CheckOverAndRewardsDefence();
                break;

            case Targets.Stay:
                // Stay
                // endless
                nowReward = 0;
                endReward = 0;
                endTypeInt = (int)EndType.Running;
                break;

            default:
                //free kill
                (endTypeInt, nowReward, endReward) = CheckOverAndRewardsFreeKill();
                break;
        }
        envUICon.ShowResult(endTypeInt);
        return (endTypeInt, nowReward, endReward);
    }

    private (int, float, float) CheckOverAndRewardsGo()
    {
        int endTypeInt = 0;
        float nowReward = 0;
        float endReward = 0;
        float nowDistance = 0;

        (nowDistance, targetCon.inArea) = sceneBlockCon.GetAgentTargetDistanceAndInside(agentObj.transform.position);
        envUICon.UpdateTargetGauge(sceneBlockCon.nowBlock.firebasesBelong, sceneBlockCon.nowBlock.belongMaxPoint);
        float areaTargetReward = GetDistanceReward(nowDistance, targetCon.inArea);
        //if(inArea != 0)
        if (sceneBlockCon.nowBlock.firebasesBelong >= sceneBlockCon.nowBlock.belongMaxPoint)
        {
            // win
            // let the area belongs to me
            nowReward = areaTargetReward;
            endReward = paramCon.goWinReward;
            //nowReward = (paramCon.inAreaReward * inArea) + getDistanceReward(nowDistance);
            endTypeInt = (int)EndType.Win;
        }
        else if (targetCon.leftTime <= 0)
        {
            // time out lose
            nowReward = areaTargetReward;
            endReward = commonParamCon.loseReward;
            endTypeInt = (int)EndType.Lose;
        }
        else
        {
            // keep on keeping on!
            nowReward = areaTargetReward;
            endReward = 0;
            endTypeInt = (int)EndType.Running;
        }
        return (endTypeInt, nowReward, endReward);
    }

    private (int, float, float) CheckOverAndRewardsAttack()
    {
        int endTypeInt = 0;
        float nowReward = 0;
        float endReward = 0;
        float nowDistance = 0;
        (nowDistance, targetCon.inArea) = sceneBlockCon.GetAgentTargetDistanceAndInside(agentObj.transform.position);
        envUICon.UpdateTargetGauge(sceneBlockCon.nowBlock.firebasesBelong, sceneBlockCon.nowBlock.belongMaxPoint);
        if (sceneBlockCon.nowBlock.GetInAreaNumber(commonParamCon.group2Tag) <= 0 && targetCon.targetEnemySpawnFinish)
        {
            // win
            // let the area belongs to me and kill every enmy in this area.
            nowReward = 0;
            endReward = paramCon.attackWinReward;
            //nowReward = (paramCon.inAreaReward * inArea) + getSceneReward(nowDistance);
            endTypeInt = (int)EndType.Win;
            targetCon.targetEnemySpawnFinish = false;
        }
        else if (targetCon.leftTime <= 0 && targetCon.targetEnemySpawnFinish)
        {
            // time out lose
            nowReward = 0;
            endReward = commonParamCon.loseReward;
            //nowReward = (paramCon.inAreaReward * inArea) + getSceneReward(nowDistance);
            endTypeInt = (int)EndType.Lose;
            targetCon.targetEnemySpawnFinish = false;
        }
        else
        {
            // keep on keeping on!
            // nowReward = (paramCon.inAreaReward * inArea) + getDistanceReward(nowDistance);
            nowReward = 0;
            endReward = 0;
            targetCon.targetEnemySpawnFinish = true;
            endTypeInt = (int)EndType.Running;
        }
        return (endTypeInt, nowReward, endReward);
    }

    private (int, float, float) CheckOverAndRewardsDefence()
    {
        // !!! NOT FINISHED YET!!!
        int endTypeInt = 0;
        float nowReward = 0;
        float endReward = 0;
        float nowDistance = 0;
        (nowDistance, targetCon.inArea) = sceneBlockCon.GetAgentTargetDistanceAndInside(agentObj.transform.position);
        envUICon.UpdateTargetGauge(sceneBlockCon.nowBlock.firebasesBelong, sceneBlockCon.nowBlock.belongMaxPoint);
        if (targetCon.leftTime <= 0 && sceneBlockCon.nowBlock.firebasesBelong >= 0f)
        {
            // win
            // time over and the area still mine
            nowReward = paramCon.defenceWinReward;
            //nowReward = (paramCon.inAreaReward * inArea) + getSceneReward(nowDistance);
            endTypeInt = (int)EndType.Win;
        }
        else if (sceneBlockCon.nowBlock.firebasesBelong <= sceneBlockCon.nowBlock.belongMaxPoint)
        {
            // lost area lose
            nowReward = commonParamCon.loseReward;
            //nowReward = (paramCon.inAreaReward * inArea) + getSceneReward(nowDistance);
            endTypeInt = (int)EndType.Lose;
        }
        else
        {
            // keep on keeping on!
            // nowReward = (paramCon.inAreaReward * inArea) + getDistanceReward(nowDistance);
            endTypeInt = (int)EndType.Running;
        }
        return (endTypeInt, nowReward, endReward);
    }

    private (int, float, float) CheckOverAndRewardsFreeKill()
    {
        int endTypeInt = 0;
        float nowReward = 0;
        float endReward = 0;
        if (enemyContainerObj.transform.childCount <= 0)
        {
            // win
            // nowReward = paramCon.winReward + (paramCon.timeBonusPerSecReward * leftTime);
            nowReward = 0;
            endReward = paramCon.freeWinReward;
            endTypeInt = (int)EndType.Win;
        }
        else if (targetCon.leftTime <= 0)
        {
            // lose
            //nowReward = paramCon.loseReward;
            nowReward = 0;
            endReward = commonParamCon.loseReward;
            endTypeInt = (int)EndType.Lose;
        }
        else
        {
            // keep on keeping on!
            nowReward = 0;
            endReward = 0;
            endTypeInt = (int)EndType.Running;
        }
        return (endTypeInt, nowReward, endReward);
    }

    /// <summary>
    /// Calculates scene reward based on distance, granting higher rewards for being closer to the target.
    /// </summary>
    /// <param name="nowDistance">The current distance.</param>
    /// <param name="inarea">Whether inside an area.</param>
    /// <returns>The reward value calculated based on distance.</returns>
    private float GetDistanceReward(float nowDistance, int inarea)
    {
        if (firstRewardFlag)
        {
            // first distance record
            (lastDistance, _) = sceneBlockCon.GetAgentTargetDistanceAndInside(agentObj.transform.position);
            firstRewardFlag = false;
        }
        float nowSeneReward = 0f;
        if (inarea != 0)
        {
            // in area
            nowSeneReward = paramCon.inAreaReward;
        }
        else
        {
            // out of area
            // nowSeneReward = paramCon.distanceReward * Math.Clamp(lastDistance - nowDistance, 0, 100);
            nowSeneReward = commonParamCon.distanceReward * (lastDistance - nowDistance);
        }
        lastDistance = nowDistance;
        return nowSeneReward;
    }

    /// <summary>
    /// Calculates kill reward based on the position of the killed enemy.
    /// </summary>
    /// <param name="enemyPosition">The position of the killed enemy.</param>
    /// <returns>The reward value calculated based on the kill position.</returns>
    public float KillReward(Vector3 enemyPosition)
    {
        float nowKillReward = 0f;
        if (targetCon.targetType == Targets.Attack)
        {
            // attack mode
            (_, int isInArea) = sceneBlockCon.nowBlock.GetDistInArea(enemyPosition);
            if (isInArea == 1)
            {
                // kill in area enemy
                nowKillReward = paramCon.killTargetEnemyReward;
            }
            else
            {
                nowKillReward = commonParamCon.killNonTargetReward;
            }
        }
        else if (targetCon.targetType == Targets.Free)
        {
            // free mode hit
            nowKillReward = paramCon.killTargetEnemyReward;
        }
        else
        {
            // goto & defence
            nowKillReward = commonParamCon.killNonTargetReward;
        }
        return nowKillReward;
    }

    /// <summary>
    /// Calculates hit reward based on the position of the hit enemy and the current mode.
    /// </summary>
    /// <param name="enemyPosition">The position of the hit enemy.</param>
    /// <returns>The reward value calculated based on the hit position and mode.</returns>
    public float HitEnemyReward(Vector3 enemyPosition)
    {
        float nowHitReward = 0f;
        if (targetCon.targetType == Targets.Attack)
        {
            // attack mode
            (_, int isInArea) = sceneBlockCon.nowBlock.GetDistInArea(enemyPosition);
            if (isInArea == 1)
            {
                // hit in area enemy
                nowHitReward = paramCon.hitTargetReward;
            }
            else
            {
                // hit not in area enemy
                nowHitReward = commonParamCon.hitNonTargetReward;
            }
        }
        else if (targetCon.targetType == Targets.Free)
        {
            // free mode hit
            nowHitReward = paramCon.hitTargetReward;
        }
        else
        {
            // goto & defence
            nowHitReward = commonParamCon.hitNonTargetReward;
        }
        return nowHitReward;
    }
}