Source code for verigym.policy.policy

from typing import Any
from abc import abstractmethod

from collections import defaultdict

from ..abstraction.abstractionmapper import AbstractionMapper
# import gymnasium as gym



class PolicyClass:
    """
    An abstract class providing an interface to any kind of model policy.
    Given an observation in the VeriGym environment,
    return an action for the VeriGym environment based on the abstract model's policy.
    """

    def __init__(
        self, policy: Any, abstraction_mapper: AbstractionMapper
    ):
        """
        Initializes a policy.

        Parameters
        ----------
        policy : Any
            The object that contains the policy (e.g. computed by an external framework). Can be used in the `self._action_from_policy
        abstraction_mapper : AbstractionMapper, optional
            An optional mapping that translates actions from e.g. abstracted environment to original environment. By default None, then an identity mapping is used, not changing the outputted action.
        """
        self.policy = policy
        self.abstraction_mapper = abstraction_mapper

    @abstractmethod
    def _action_from_policy(self, obs):
        """
        Get an action from the model's policy by providing an observation/state that from the abstract space.

        Parameters
        ----------
        obs : object
            an observation in the abstract space

        Returns
        -------
        action : object
            an (abstract) action according to the policy
        """
        # This should be implemented in specific child classes
        raise NotImplementedError



    def get_action(self, obs, info=None):
        """
        Gets an observation from `env.observation_space`.
        Queries the model's policy.
        Map the action from the model's policy to `env.action_space`

        Parameters
        ----------
        obs : object
            An observation in the environment's observation space

        Returns
        -------
        action : object
            An action in the environment's action space

        """
        o = self.abstraction_mapper.original_to_abstract_state(
            obs
        )  # self._obs_to_model(obs)
        a = self._action_from_policy(o)
        action = self.abstraction_mapper.abstract_to_original_action(a)
        return action




    def seed(self, seed: int | None = None) -> None:
        """
        Seeds the random generator of the policy, e.g. by `create_abstraction(seed=...)`.
        Deterministic policies ignore it; policies that sample their actions need to redefine it.

        Parameters
        ----------
        seed : int | None, optional
            The seed. If None, the generator is seeded from OS entropy.
        """
        pass




    def update_for_abstraction_refinement(
        self,
        dataset: list[tuple],
        T_counts: defaultdict[int, defaultdict[int, defaultdict[int, int]]],
        P_tot: dict,
        R_cumsum: defaultdict[int, defaultdict[int, float]],
        S_init_counts: defaultdict[int, int],
    ) -> "PolicyClass":
        """
        Updates the policy during the exploration phase of the abstraction learning phase.
        - If this policy does not change during abstraction learning, it just returns itself, unchanged.
        - If the policy does update, this function needs to be redefined for that policy class.

        Examples for which policies this function may be useful:
        - state-based sampling (e.g. entropy-based sampling)
        - RL policy that is updated throughout the abstraction learning process.

        Parameters
        ----------
        dataset: list[tuple]
            A dataset with observations, list of tuples (state, action, reward, next_state).
        T_counts : defaultdict[int, defaultdict[int, defaultdict[int, int]]]
            A nested dict `T_count[state][action][next_state]` returns the integer count.
        P_tot : dict[tuple[int, int], int]
            A dict storing the occurrences of state-action pairs. Useful for iterating through `T_counts` and `R_cumsum`, see `learn_abstraction.py` for examples.
        R_cumsum : defaultdict[int, defaultdict[int, float]]
            A nested dict, where `R_cumsum[state][action]` returns the cumulative reward encountered for the given state-action pair (not summed up over time).
            The mean reward is `R_cumsum[state][action] / P_tot[(state, action)]`.
        S_init_counts : defaultdict[int, int]
            A dict, where `S_init_counts[state]` returns the number of occurences that `state` was the intial state.

        Returns
        -------
        PolicyClass
            The policy that should be further used for exploring through the abstraction learning.
        """
        return self