In __init__, iterating over FinitProbabilityDistribution simplified

Author: coverdrive

rl/markov_decision_process.py

@@ -134,7 +134,7 @@ def __init__(
         non_terminals: Set[S] = set(mapping.keys())
         self.mapping = {NonTerminal(s): {a: Categorical(
             {(NonTerminal(s1) if s1 in non_terminals else Terminal(s1), r): p
-             for (s1, r), p in v.table().items()}
+             for (s1, r), p in v}
         ) for a, v in d.items()} for s, d in mapping.items()}
         self.non_terminal_states = list(self.mapping.keys())
 
@@ -165,7 +165,7 @@ def apply_finite_policy(self, policy: FinitePolicy[S, A])\
                 = defaultdict(float)
             actions = policy.act(state)
             for action, p_action in actions:
-                for (s1, r), p in action_map[action].table().items():
+                for (s1, r), p in action_map[action]:
                     outcomes[(s1.state, r)] += p_action * p
 
             transition_mapping[state.state] = Categorical(outcomes)

rl/markov_process.py

@@ -104,7 +104,7 @@ def __init__(self, transition_map: Mapping[S, FiniteDistribution[S]]):
         self.transition_map = {
             NonTerminal(s): Categorical(
                 {(NonTerminal(s1) if s1 in non_terminals else Terminal(s1)): p
-                 for s1, p in v.table().items()}
+                 for s1, p in v}
             ) for s, v in transition_map.items()
         }
         self.non_terminal_states = list(self.transition_map.keys())
@@ -272,7 +272,7 @@ def __init__(
         self.transition_reward_map = {
             NonTerminal(s): Categorical(
                 {(NonTerminal(s1) if s1 in nt else Terminal(s1), r): p
-                 for (s1, r), p in v.table().items()}
+                 for (s1, r), p in v}
             ) for s, v in transition_reward_map.items()
         }