PrimAITE/src/primaite/game/agent/rewards.py

from abc import ABC, abstractmethod
from typing import Any, Dict, List


class AbstractReward:
    def __init__(self):
        ...

    @abstractmethod
    def calculate(self, state: Dict) -> float:
        return 0.3


class DummyReward(AbstractReward):
    def calculate(self, state: Dict) -> float:
        return -0.1


class RewardFunction:
    __rew_class_identifiers: Dict[str, type[AbstractReward]] = {"DUMMY": DummyReward}

    def __init__(self, reward_function: AbstractReward):
        self.reward: AbstractReward = reward_function

    def calculate(self, state: Dict) -> float:
        return self.reward.calculate(state)

    @classmethod
    def from_config(cls, cfg: Dict) -> "RewardFunction":
        for rew_component_cfg in cfg["reward_components"]:
            rew_type = rew_component_cfg["type"]
            rew_component = cls.__rew_class_identifiers[rew_type]()
            new = cls(reward_function=rew_component)
        return new