<reponame>nicoguertler/graphrl<gh_stars>1-10
from abc import ABC, abstractmethod

from .subtask_specs import SubtaskSpec
from ..spaces import space_from_gym_space


class ReturnMaximSubtaskSpec(SubtaskSpec):
    """Specification of return maximization subtask, i.e., obs, reward etc."""

    def __init__(self, max_n_actions=None):
        super().__init__()
        self._max_n_actions = max_n_actions

    @property
    @abstractmethod
    def obs_space(self):
        """Observation space."""

        pass

    @abstractmethod
    def map_to_obs(self, env_obs, parent_info, ep_step):
        """Maps environment observation, parent info and episode step to observation."""

        pass

    @abstractmethod
    def get_reward(self, obs, act, new_obs, parent_info, ep_step, env_info):
        """Get reward of transition."""

        pass

    @abstractmethod
    def step_update(self, env_info, parent_info, sess_info):
        """Update internal state of subtask spec in each time step."""

        pass

    def reset(self):
        """Reset state of subtask sepc."""

        pass

    @property
    def max_n_actions(self):
        """Maximum number of actions before subtask ends."""

        return self._max_n_actions


class EnvRMSubtaskSpec(ReturnMaximSubtaskSpec):
    """Return maximization subtask specification based on environment."""

    def __init__(self, env, max_n_actions=None):
        super().__init__(max_n_actions)
        self._env = env
        self._env_obs_space = space_from_gym_space(env.observation_space)
        self._cum_reward = 0.

    @property
    def obs_space(self):
        """Use environment observation space."""

        return self._env_obs_space

    def map_to_obs(self, env_obs, parent_info, ep_step):
        """Return environment observation."""

        return env_obs

    def get_reward(self, obs, act, new_obs, parent_info, ep_step, env_info):
        """Return cumulative environment reward.
        
        Reward has been accumulated since las call to get_reward 
        or beginning of episode.
        """

        rew = self._cum_reward + env_info.reward
        self._cum_reward = 0

        return rew

    def step_update(self, env_info, parent_info, sess_info):
        """Add env reward of time step to cumulative reward."""

        self._cum_reward += env_info.reward

from abc import ABC, abstractmethod

from .subtask_specs import SubtaskSpec
from ..spaces import space_from_gym_space


class ReturnMaximSubtaskSpec(SubtaskSpec):
    """Specification of return maximization subtask, i.e., obs, reward etc."""

    def __init__(self, max_n_actions=None):
        super().__init__()
        self._max_n_actions = max_n_actions

    @property
    @abstractmethod
    def obs_space(self):
        """Observation space."""

        pass

    @abstractmethod
    def map_to_obs(self, env_obs, parent_info, ep_step):
        """Maps environment observation, parent info and episode step to observation."""

        pass

    @abstractmethod
    def get_reward(self, obs, act, new_obs, parent_info, ep_step, env_info):
        """Get reward of transition."""

        pass

    @abstractmethod
    def step_update(self, env_info, parent_info, sess_info):
        """Update internal state of subtask spec in each time step."""

        pass

    def reset(self):
        """Reset state of subtask sepc."""

        pass

    @property
    def max_n_actions(self):
        """Maximum number of actions before subtask ends."""

        return self._max_n_actions


class EnvRMSubtaskSpec(ReturnMaximSubtaskSpec):
    """Return maximization subtask specification based on environment."""

    def __init__(self, env, max_n_actions=None):
        super().__init__(max_n_actions)
        self._env = env
        self._env_obs_space = space_from_gym_space(env.observation_space)
        self._cum_reward = 0.

    @property
    def obs_space(self):
        """Use environment observation space."""

        return self._env_obs_space

    def map_to_obs(self, env_obs, parent_info, ep_step):
        """Return environment observation."""

        return env_obs

    def get_reward(self, obs, act, new_obs, parent_info, ep_step, env_info):
        """Return cumulative environment reward.
        
        Reward has been accumulated since las call to get_reward 
        or beginning of episode.
        """

        rew = self._cum_reward + env_info.reward
        self._cum_reward = 0

        return rew

    def step_update(self, env_info, parent_info, sess_info):
        """Add env reward of time step to cumulative reward."""

        self._cum_reward += env_info.reward
Specification of return maximization subtask, i.e., obs, reward etc. Observation space. Maps environment observation, parent info and episode step to observation. Get reward of transition. Update internal state of subtask spec in each time step. Reset state of subtask sepc. Maximum number of actions before subtask ends. Return maximization subtask specification based on environment. Use environment observation space. Return environment observation. Return cumulative environment reward.                  Reward has been accumulated since las call to get_reward          or beginning of episode. Add env reward of time step to cumulative reward.