colosseum.utils.loops

  1from typing import TYPE_CHECKING, Iterable, List, Tuple, Union
  2
  3import numpy as np
  4from dm_env import TimeStep
  5
  6from colosseum import config
  7
  8if TYPE_CHECKING:
  9    from colosseum.mdp import ContinuousMDP, EpisodicMDP
 10
 11
 12def human_loop(mdp: Union["ContinuousMDP", "EpisodicMDP"], other_policies: dict = None):
 13    """
 14    allows a human to control an MDP.
 15    """
 16
 17    verba = lambda: print(mdp.get_grid_representation(mdp.cur_node))
 18
 19    print("Start calculating the optimal policy")
 20    optimal_policy = mdp.get_optimal_policy(False)
 21    print("End calculating the optimal policy")
 22
 23    state = mdp.reset()
 24    while True:
 25        print("State:", state)
 26        verba()
 27
 28        if mdp.is_episodic():
 29            optimal_action = optimal_policy[mdp.h, mdp.node_to_index[mdp.cur_node]]
 30        else:
 31            optimal_action = optimal_policy[mdp.node_to_index[mdp.cur_node]]
 32        print(f"The optimal action for this state is:{optimal_action}")
 33
 34        if other_policies is not None:
 35            for pi_name, pi in other_policies.items():
 36                print(
 37                    f"The action of policy {pi_name} for this state is:{np.argmax(pi[mdp.cur_node])}"
 38                )
 39
 40        action = int(
 41            input(
 42                "Available actions are: "
 43                + ",".join(map(str, range(mdp.n_actions)))
 44                + ".\tChoose one to act or type anything else to terminate.\n"
 45            )
 46        )
 47        if action not in range(mdp.n_actions):
 48            break
 49        state = mdp.step(action)
 50        if state.last():
 51            print("State:", state)
 52            state = mdp.reset()
 53
 54
 55def random_loop(
 56    mdp: Union["ContinuousMDP", "EpisodicMDP"],
 57    N: int,
 58    return_actions: bool = False,
 59    human_readable=False,
 60) -> Union[Tuple[List[TimeStep], List[int]], List[TimeStep]]:
 61    """
 62    generates interactions data by selecting actions a random.
 63
 64    Parameters
 65    ----------
 66    mdp: Union["ContinuousMDP", "EpisodicMDP"]
 67        The MDP instance.
 68    N : int
 69        The number of interactions.
 70    return_actions: bool, optional
 71        If True, the selected actions are returned. By default, it is set to False,
 72    human_readable: bool
 73        If True, the state information is printed in a human interpretable form. By default, it is set to False.
 74
 75    Returns
 76    -------
 77    Union[Tuple[List[TimeStep], List[int]], List[TimeStep]]
 78        The data generated from the interactions.
 79    """
 80
 81    if human_readable:
 82        verba = lambda: print(mdp.get_grid_repr())
 83    else:
 84        verba = lambda: print("State:", state, "Action: ", action)
 85
 86    states = []
 87    state = mdp.reset()
 88    states.append(state)
 89    actions = []
 90    action = None
 91    while len(states) < N:
 92        if config.VERBOSE_LEVEL > 0:
 93            verba()
 94        state, action = mdp.random_step()
 95        if return_actions:
 96            actions.append(action)
 97        states.append(state)
 98        if state.last():
 99            if config.VERBOSE_LEVEL > 0:
100                print("Last state:", state)
101            state = mdp.reset()
102            states.append(state)
103
104    if return_actions:
105        return states, actions
106    return states
107
108
109def prefixed_action_loop(
110    mdp: Union["ContinuousMDP", "EpisodicMDP"],
111    actions: Iterable[int],
112    human_readable: bool = False,
113) -> List[TimeStep]:
114    """
115    generates interaction with the MDP according to the actions given in input.
116
117    Parameters
118    ----------
119    mdp: Union["ContinuousMDP", "EpisodicMDP"]
120        The MDP instance.
121    actions : Iterable[int]
122        The actions to be selected.
123    human_readable: bool
124        If True, the state information is printed in a human interpretable form. By default, it is set to False.
125
126    Returns
127    -------
128    Union[Tuple[List[TimeStep], List[int]], List[TimeStep]]
129        The data generated from the interactions.
130    """
131
132    if human_readable:
133        verba = lambda: print(mdp.get_grid_repr())
134    else:
135        verba = lambda: print("State:", state, "Action: ", action)
136
137    states = []
138    state = mdp.reset()
139    states.append(state)
140    for action in actions:
141        if config.VERBOSE_LEVEL > 0:
142            verba()
143        state = mdp.step(action)
144        states.append(state)
145        if state.last():
146            if config.VERBOSE_LEVEL > 0:
147                print("Last state:", state)
148            state = mdp.reset()
149            states.append(state)
150    return states