colosseum.dynamic_programming.utils

  1import numba
  2import numpy as np
  3
  4ARGMAX_SEED = 42
  5_rng = np.random.RandomState(ARGMAX_SEED)
  6
  7
  8class DynamicProgrammingMaxIterationExceeded(Exception):
  9    pass
 10
 11
 12@numba.njit()
 13def argmax_2d(A: np.ndarray) -> np.ndarray:
 14    """
 15    Returns
 16    -------
 17    np.ndarray
 18        The array of same dimensionality of A with ones corresponding to the max across rows and zeros otherwise.
 19    """
 20    np.random.seed(ARGMAX_SEED)
 21    X = np.zeros_like(A, np.float32)
 22    for s in range(len(A)):
 23        i = np.random.choice(np.where(A[s] == A[s].max())[0])
 24        X[s, i] = 1
 25    return X
 26
 27
 28@numba.njit()
 29def argmax_3d(A: np.ndarray) -> np.ndarray:
 30    """
 31    implements a vectorized version of `argmax_2d`.
 32    """
 33    np.random.seed(ARGMAX_SEED)
 34    X = np.zeros(A.shape, np.float32)
 35    for h in range(len(A)):
 36        for s in range(A.shape[1]):
 37            i = np.random.choice(np.where(A[h, s] == A[h, s].max())[0])
 38            X[h, s, i] = 1.0
 39    return X
 40
 41
 42@numba.njit()
 43def get_deterministic_policy_from_q_values(Q: np.ndarray) -> np.ndarray:
 44    """
 45    Returns
 46    -------
 47    np.ndarray
 48        The infinite horizon optimal deterministic policy for the given q values.
 49    """
 50    np.random.seed(ARGMAX_SEED)
 51    X = np.zeros(Q.shape[:-1], np.int32)
 52    for s in range(len(Q)):
 53        i = np.random.choice(np.where(Q[s] == Q[s].max())[0])
 54        X[s] = np.int32(i)
 55    return X
 56
 57
 58@numba.njit()
 59def get_deterministic_policy_from_q_values_finite_horizon(Q: np.ndarray) -> np.ndarray:
 60    """
 61    Returns
 62    -------
 63    np.ndarray
 64        The finite horizon optimal deterministic policy for the given q values.
 65    """
 66    np.random.seed(ARGMAX_SEED)
 67    X = np.zeros(Q.shape[:-1], np.int32)
 68    for h in range(len(Q)):
 69        for s in range(Q.shape[1]):
 70            i = np.random.choice(np.where(Q[h, s] == Q[h, s].max())[0])
 71            X[h, s] = np.int32(i)
 72    return X
 73
 74
 75def get_policy_from_q_values(Q: np.ndarray, stochastic_form=False) -> np.ndarray:
 76    """
 77    Parameters
 78    ----------
 79    Q : np.ndarray
 80        The q-value estimates.
 81    stochastic_form : bool
 82        If False, the array contains the integers corresponding to the optimal actions. If True, the array contains
 83        vectors representing the deterministic probability distributions.
 84
 85    Returns
 86    -------
 87    np.ndarray
 88        The deterministic policy derived from the q_values given in input.
 89
 90    """
 91    # Episodic case
 92    if Q.ndim == 3:
 93        if stochastic_form:
 94            return argmax_3d(Q)
 95        return get_deterministic_policy_from_q_values_finite_horizon(Q)
 96
 97    # Infinite horizon case
 98    if stochastic_form:
 99        return argmax_2d(Q)
100    return get_deterministic_policy_from_q_values(Q)