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