AINetworkDefense_Competitive/rewards.py

import tensorflow as tf

def compute_rewards(mode, env_info):
    """
    Calculates zero-sum competitive rewards based on live environment info keys.
    Guarantees a 2-value tuple output to prevent unpacking errors.
    """
    if not env_info:
        return 0.0, 0.0

    score_diff = env_info.get('score_difference', 0.0)
    progress = env_info.get('team_progress', 0.0)
        
    # Extract values safely regardless of whether they arrive as scalars or tensors
    if hasattr(score_diff, 'numpy'): score_diff = score_diff.numpy()
    if hasattr(progress, 'numpy'): progress = progress.numpy()
        
    score_diff = float(score_diff)
    progress = float(progress)

    if str(mode).lower() == "competitive":
        # Red gets rewarded when score difference grows, but pays a step penalty
        red_reward = score_diff + (progress * 0.1) - 0.5            
        # Blue gets rewarded when score difference shrinks, plus a survival bonus
        blue_reward = -score_diff - (progress * 0.1) + 0.5
    else:
        # Cooperative or fallback layout
        red_reward = score_diff
        blue_reward = progress

    return float(red_reward), float(blue_reward)