Examples

Environment

Before implementing the reinforcement learning algorithm, it is very useful to implement the real world from which episodes will be sampled. This environment will be based on a known MDP, and its job is to yield on demand random episodes which respect asymptotically the MDP (transitions and rewards). The goal of this entire section is to implement, test the reinforcement learning algorithms and show that they enable to learn from the interactions with the environment. Depending on the learning algorithm, either the MDP itself is estimated, or the qualities of a policy, or the optimum policy.

src

"""Base types for a sampling environment"""

from collections.abc import Generator
from typing import Protocol
from collections import namedtuple
from dataclasses import dataclass, field

from ai.reinforcement.mdp_types import (
    State,
    Action,
    Transition,
    Reward,
    Policy,
)


type Outcome[S] = tuple[S, Reward]  # careful, singular, single outcome

Frame = namedtuple("Frame", ("state", "action", "reward", "next_state"))


@dataclass
class OutcomesForChoice[S: State]:  # only reformatting mdp to ease random choice
    outcomes: list[Outcome[S]] = field(default_factory=list)
    transitions: list[Transition] = field(default_factory=list)


class Environment[S: State](Protocol):
    """an environnement is able to spit frames of episodes on demand.
    A frame is a ("state", "action", "reward", "next_state") named tuple.
    """

    def actions_from_state(self, state: S) -> set[Action]: ...

    def outcome_from_state_and_action(self, state: S, action: Action) -> Outcome[S]: ...

    """careful, return a single random outcome following a given probability distribution"""

    def is_end_state(self, state: S) -> bool: ...

    def yield_frame_in_episode(self, state: S, policy: Policy) -> Generator[Frame]: ...

Exercise (intermediate)

Implement and test EnvironmentFromMDP.

Correction

src

class EnvironmentFromMDP(Environment):
    def __init__(self, mdp: MDPProblem):
        self.mdp = mdp
        self.outcomes_for_choice = self.format_outcomes_for_choice()

    def format_outcomes_for_choice(
        self,
    ) -> dict[tuple[State, Action], OutcomesForChoice]:
        transition_cases = {}
        for state in self.mdp.states:
            for action in self.mdp.actions_from_state(state):
                outcomes: Outcomes = self.mdp.outcomes_from_state_and_action(
                    state, action
                )
                outcomes_for_choice = OutcomesForChoice()
                for next_state, (transition, reward) in outcomes.items():
                    outcomes_for_choice.outcomes.append((next_state, reward))
                    outcomes_for_choice.transitions.append(transition)

                transition_cases[(state, action)] = outcomes_for_choice
        return transition_cases

    @override
    def actions_from_state(self, state: State) -> set[Action]:
        return self.mdp.actions_from_state(state)

    @override
    def outcome_from_state_and_action(self, state: State, action: Action) -> Outcome:
        return choices(
            self.outcomes_for_choice[(state, action)].outcomes,
            self.outcomes_for_choice[(state, action)].transitions,
        )[0]

    @override
    def is_end_state(self, state: State) -> bool:
        return not self.actions_from_state(state)

    @override
    def yield_frame_in_episode(self, state: State, policy: Policy) -> Generator[Frame]:
        while not self.is_end_state(state):
            actions = self.actions_from_state(state)
            action = policy.action_from_state(state, actions)
            next_state, reward = self.outcome_from_state_and_action(state, action)
            yield Frame(state, action, reward, next_state)
            state = next_state

tests

import pytest

from ai.reinforcement.environment_types import OutcomesForChoice


def test_format_outcomes_for_choice(game_environment):
    expected = {
        ("in", "quit"): OutcomesForChoice([("out", 10.0)], [1.0]),
        ("in", "stay"): OutcomesForChoice(
            [("in", 4.0), ("out", 5.0)], [2.0 / 3.0, 1.0 / 3.0]
        ),
    }
    assert expected == game_environment.outcomes_for_choice  # dubious float comparison


def test_outcome_from_state_and_action(game_environment):
    nb_trials = 1_000_000
    nb_in = 0
    for _ in range(nb_trials):
        next_state, _ = game_environment.outcome_from_state_and_action("in", "stay")
        if next_state == "in":
            nb_in += 1
    assert pytest.approx(nb_in / nb_trials, abs=0.01) == 2.0 / 3.0


def test_yield_frame_in_episode(game_environment, policy_stay):
    nb_episode = 100_000
    nb_in = 0
    nb_frame = 0
    for _ in range(nb_episode):
        for frame in game_environment.yield_frame_in_episode("in", policy_stay):
            nb_frame += 1
            if frame.next_state == "in":
                nb_in += 1
    assert pytest.approx(nb_in / nb_frame, abs=0.03) == 2.0 / 3.0

Game

Exercise (difficult)

For the four following questions, use a RandomPolicy.

1. Implement and test the model based Monte-Carlo algorithm. Add a export_to_MDPMapping method.

2. Implement and test the model free Monte-Carlo algorithm.

3. Implement and test the Sarsa algorithm.

4. Implement and test the Q-learning algorithm.

Hints

The four algorithms work iteratively: the estimation is incrementally improved. The easiest way is to loop a certain amount of times over all discovered states, each time generate an episode and use it with one of the algorithm.

Obviously, it would be even better have a convergence check: either for every estimated transition, or quality. This step is not implemented in the correction.

These algorithms are not easy to test because of the randomness from the environment sampling. However, if the environment is an input parameter of the function, it is a prime example of a dependency injection. It allows to feed a tweaked environment which behaves the same, but which is not random. For example:

tests

class DummyGameEnvironment(EnvironmentFromMDP):
    def __init__(self, mdp: MDPProblem):
        super().__init__(mdp)
        self.outcome_cycle = cycle((("in", 4.0), ("in", 4.0), ("out", 5.0)))

    def outcome_from_state_and_action(self, state: State, action: Action) -> Outcome:
        return next(self.outcome_cycle)

This environment will always produce the same episode.

Hints 1.

Maintain the number of times a (state, action, next_state) and a (state, action) has been encountered. The ratio between the two is the estimation of a transition. Use the class StateActionNextStateMetric in a defaultdict.

src

@dataclass
class StateActionNextStateMetric:
    counter: int = 0
    reward_sum: Reward = 0.0

Do not forget to update your knowledge of the existing states.

An approximation of the true transitions is expected.

Hints 2.

Maintain the number of times (state, action) has been encountered because it will be the \(N\) in the convex combination. Use reversed full episode and maintain the qualities for every (state, action). Going in reverse, allows to easily calculate the discounted utility: \(u_t=r+\gamma u_{t+1}\). If you update the quality (state, action) for every frame, it will automatically keep the first encounter of every (state, action) of the episode.

It ease much easier to first obtain the utilities for every (state, action) of the episode, and then perform the convex combination to update the qualities.

An approximation of the random policy qualities are expected.

Hints 3.

This the hardest of the four to implement because you need a pick view of the next action. However, a frame is state, action, reward, next_state. At a current frame, you do not have the next action. A possibility is to loop by hand with a while loop, calling next on the episode iterator. This way you can manage to constantly have access to a frame and its successor. However, pay attention if the next state is an end state or not.

An approximation of the random policy qualities are expected.

Hints 4.

Similar to the previous one, but easier. You no longer need the next action, just the next optimal action from the next state. Do not forget to update the optimal policy mapping right after the convex combination update.

An approximation of the optimal policy qualities are expected.

Correction

src

def model_based_Monte_Carlo(
    environment: Environment,
    state_: State,
    policy: Policy,
    nb_episodes_per_state: int = 1_000,
) -> MDPMapping:
    """builds an estimated MDP and returns its mapping"""
    states: set[State] = {state_}

    S_A_NS_metric: dict[tuple[State, Action, NextState], StateActionNextStateMetric] = (
        defaultdict(StateActionNextStateMetric)
    )  # numerator transition fraction and denominator reward fraction
    counter_state_action: dict[tuple[State, Action], int] = defaultdict(int)
    # denominator transition fraction

    # generate episode many times from all known states
    for _ in range(nb_episodes_per_state):
        for state in states.copy():  # must not iterate on a mutating set
            episode: Generator[Frame] = environment.yield_frame_in_episode(
                state, policy
            )

            if environment.is_end_state(state):
                continue

            # analyze episode
            for state, action, reward, next_state in episode:
                states.add(state)
                states.add(next_state)
                counter_state_action[(state, action)] += 1
                metric = S_A_NS_metric[(state, action, next_state)]
                metric.counter += 1
                metric.reward_sum += reward
    return export_to_MDPMapping(states, S_A_NS_metric, counter_state_action)

src

def export_to_MDPMapping(
    states: set[State],
    S_A_NS_metric: dict[tuple[State, Action, NextState], StateActionNextStateMetric],
    counter_state_action: dict[tuple[State, Action], int],
) -> MDPMapping:
    # [state][action][next_state] = (Transition, Reward)
    mdp: MDPMapping = {}
    for state in states:
        mdp[state] = {}
    for state, action in counter_state_action:
        mdp[state][action] = {}

    for (state, action, next_state), metric in S_A_NS_metric.items():
        mdp[state][action][next_state] = (
            metric.counter / counter_state_action[(state, action)],
            metric.reward_sum / metric.counter,
        )
    return mdp

tests

def test_model_based_Monte_Carlo_dummy(dummy_game_environment, policy_stay):
    mdp_mapping: MDPMapping = model_based_Monte_Carlo(
        dummy_game_environment, "in", policy_stay, 2
    )
    assert pytest.approx(mdp_mapping["in"]["stay"]["in"][0]) == 2.0 / 3.0
    assert pytest.approx(mdp_mapping["in"]["stay"]["in"][1]) == 4.0

src

def model_free_Monte_Carlo(
    environment: Environment,
    state_: State,
    policy: Policy,
    discount: float = 1.0,
    nb_episodes_per_state: int = 1_000,
) -> QualitiesMapping:
    states: set[State] = {state_}

    qualities: QualitiesMapping[State] = defaultdict(float)
    counter_state_action: dict[tuple[State, Action], int] = defaultdict(int)
    # counter number of update for each Q(s,a)

    # generate episode many times from all known states
    for _ in range(nb_episodes_per_state):
        for state in states.copy():  # must not iterate on a mutating set
            episode: Generator[Frame] = environment.yield_frame_in_episode(
                state, policy
            )

            if environment.is_end_state(state):
                continue

            # analyze episode
            reversed_episode = reversed(tuple(episode))
            utility: Reward = 0.0
            updated_state_action: set[tuple[State, Action]] = set()
            utilities: QualitiesMapping = {}
            # obtain first occurrence utilities
            for state, action, reward, next_state in reversed_episode:
                states.add(state)
                states.add(next_state)
                utility = utility * discount + reward
                utilities[(state, action)] = utility
                updated_state_action.add((state, action))

            # update qualities
            for state, action in updated_state_action:
                alpha: float = 1.0 / (1.0 + counter_state_action[(state, action)])
                qualities[(state, action)] = (1.0 - alpha) * qualities[
                    (state, action)
                ] + alpha * utilities[(state, action)]
                counter_state_action[(state, action)] += 1

    return qualities

tests

def test_model_free_Monte_Carlo_dummy(dummy_game_environment, policy_stay):
    qualities_mapping: QualitiesMapping = model_free_Monte_Carlo(
        dummy_game_environment, "in", policy_stay, nb_episodes_per_state=2
    )
    assert pytest.approx(qualities_mapping[("in", "stay")]) == 13.0

src

def sarsa(
    environment: Environment,
    state_: State,
    policy: Policy,
    discount: float = 1.0,
    nb_episodes_per_state: int = 1_000,
) -> QualitiesMapping:
    states: set[State] = {state_}

    qualities: QualitiesMapping = defaultdict(float)
    counter_state_action: dict[tuple[State, Action], int] = defaultdict(int)
    # counter number of update for each Q(s,a)

    # generate episode many times from all known states
    for _ in range(nb_episodes_per_state):
        for state in states.copy():  # must not iterate on a mutating set
            episode: Generator[Frame] = environment.yield_frame_in_episode(
                state, policy
            )
            if environment.is_end_state(state):
                continue

            # analyze episode
            current_frame: Frame = next(episode)
            last_frame: bool = False
            while not last_frame:
                state, action, reward, next_state = current_frame

                # obtain quality next state
                if environment.is_end_state(next_state):
                    next_quality: Reward = 0.0
                    last_frame = True
                else:
                    next_frame: Frame = next(episode)
                    _, next_action, _, _ = next_frame
                    next_quality: Reward = qualities[(next_state, next_action)]
                    current_frame = next_frame

                # update qualities
                alpha: float = 1.0 / (1.0 + counter_state_action[(state, action)])
                qualities[(state, action)] = (1.0 - alpha) * qualities[
                    (state, action)
                ] + alpha * (reward + discount * next_quality)
                counter_state_action[(state, action)] += 1

    return qualities

tests

def test_sarsa_dummy(dummy_game_environment, policy_stay):
    qualities_mapping: QualitiesMapping = sarsa(
        dummy_game_environment, "in", policy_stay, nb_episodes_per_state=1
    )
    assert pytest.approx(qualities_mapping[("in", "stay")]) == 17.0 / 3.0

src

def qlearning(
    environment: Environment,
    state_: State,
    policy: Policy,
    discount: float = 1.0,
    nb_episodes_per_state: int = 1_000,
) -> tuple[QualitiesMapping, PolicyMapping]:
    states: set[State] = {state_}

    opt_qualities: QualitiesMapping = defaultdict(float)
    counter_state_action: dict[tuple[State, Action], int] = defaultdict(int)
    # counter number of update for each Q(s,a)
    opt_policy_mapping: PolicyMapping = {}

    # generate episode many times from all known states
    for _ in range(nb_episodes_per_state):
        for state in states.copy():  # must not iterate on a mutating set
            episode: Generator[Frame] = environment.yield_frame_in_episode(
                state, policy
            )
            if environment.is_end_state(state):
                continue

            # analyze episode
            for frame in episode:
                state, action, reward, next_state = frame
                states.add(state)
                states.add(next_state)

                # obtain quality next state
                if environment.is_end_state(next_state):
                    opt_next_quality: Reward = 0.0
                else:
                    next_actions = environment.actions_from_state(next_state)
                    opt_next_quality = max(
                        [
                            opt_qualities[(next_state, next_action)]
                            for next_action in next_actions
                        ]
                    )

                # update optimum qualities
                alpha: float = 1.0 / (1.0 + counter_state_action[(state, action)])
                opt_qualities[(state, action)] = (1.0 - alpha) * opt_qualities[
                    (state, action)
                ] + alpha * (reward + discount * opt_next_quality)
                counter_state_action[(state, action)] += 1

                # update optimum policy
                actions = environment.actions_from_state(state)
                _, opt_action = max(
                    [(opt_qualities[(state, action)], action) for action in actions]
                )
                opt_policy_mapping[state] = opt_action

    return opt_qualities, opt_policy_mapping

tests

def test_qlearning_dummy(dummy_game_environment, policy_stay):
    qualities_mapping, opt_policy_mapping = qlearning(
        dummy_game_environment, "in", policy_stay, nb_episodes_per_state=1
    )
    assert pytest.approx(qualities_mapping[("in", "stay")]) == 17.0 / 3.0
    assert opt_policy_mapping["in"] == "stay"

Magic tram

Exercise (easy)

Use the four previous algorithms on the magic tram problem.

1. Model based MC, comment on the result.

2. Model free MC and Sarsa, why don’t we find the same qualities as the correction of the magic tram MDP exercise?

3. Qlearning, do we retrieve the optimal policy of the magic tram MDP exercise? Use this policy mapping in Model free MC and Sarsa to confirm that you find a satisfying estimation of the qualities of the magic tram MDP exercise.

Correction

For nb_episodes_per_state of 100_000 and a fail probability of 0.2.

1. The MDP and the fail probability are correctly estimated, the variance of the estimations between runs decreases when the nb_episodes_per_state increases.

   {2: {'tram': {2: (0.20043060556168696, -2.0), 4: (0.799569394438313, -2.0)},
         'walk': {3: (1.0, -1.0)}},
   3: {'tram': {3: (0.19742429996892588, -2.0), 6: (0.8025757000310741, -2.0)},
         'walk': {4: (1.0, -1.0)}},
   4: {'tram': {4: (0.20133871320884553, -2.0), 8: (0.7986612867911544, -2.0)},
         'walk': {5: (1.0, -1.0)}},
   5: {'tram': {5: (0.2012801210754909, -2.0), 10: (0.7987198789245091, -2.0)},
         'walk': {6: (1.0, -1.0)}},
   6: {'tram': {6: (0.20070165624923997, -2.0), 12: (0.79929834375076, -2.0)},
         'walk': {7: (1.0, -1.0)}},
   7: {'walk': {8: (1.0, -1.0)}},
   8: {'walk': {9: (1.0, -1.0)}},
   9: {'walk': {10: (1.0, -1.0)}},
   10: {'walk': {11: (1.0, -1.0)}},
   11: {'walk': {12: (1.0, -1.0)}},
   12: {'walk': {13: (1.0, -1.0)}},
   13: {}}
   

2. The magic tram MDP correction gave the qualities of the optimal policy (value iteration algorithm). Here, the random policy is used. The two matches when the two policies are equivalent: when the only choice is to walk.

   {(2, 'tram'): -9.654715245560729,
   (2, 'walk'): -9.118302645579085,
   (3, 'tram'): -7.970203666250092,
   (3, 'walk'): -8.238426328292135,
   (4, 'tram'): -7.4482404985182455,
   (4, 'walk'): -7.0281073940240875,
   (5, 'tram'): -5.600584493424383,
   (5, 'walk'): -6.443788333294008,
   (6, 'tram'): -3.896297797594899,
   (6, 'walk'): -6.999999999999999,
   (7, 'walk'): -6.0,
   (8, 'walk'): -4.999999999999999,
   (9, 'walk'): -4.0,
   (10, 'walk'): -3.0,
   (11, 'walk'): -2.0,
   (12, 'walk'): -1.0}
   
   
   {(2, 'tram'): -9.481134252634465,
   (2, 'walk'): -8.877060965638796,
   (3, 'tram'): -7.858008074723888,
   (3, 'walk'): -8.09721124858775,
   (4, 'tram'): -7.409247009049843,
   (4, 'walk'): -6.930076577752609,
   (5, 'tram'): -5.580370812949426,
   (5, 'walk'): -6.380390930193599,
   (6, 'tram'): -3.8683076395406277,
   (6, 'walk'): -6.9511901596126995,
   (7, 'walk'): -5.98235136847001,
   (8, 'walk'): -4.994789408530305,
   (9, 'walk'): -3.998836675565334,
   (10, 'walk'): -2.9998236749092415,
   (11, 'walk'): -1.9999872676343076,
   (12, 'walk'): -1.0}
   

3. Yes. Using the optimum policy from Qlearning in Sarsa or Model free MC allows to retrieve the qualities from the magic tram MDP correction.

{2: 'walk',
3: 'tram',
4: 'walk',
5: 'walk',
6: 'tram',
7: 'walk',
8: 'walk',
9: 'walk',
10: 'walk',
11: 'walk',
12: 'walk'}

{(2, 'walk'): -7.0063599999999475,
(3, 'tram'): -6.000875004375021,
(6, 'tram'): -3.5002433349554973,
(12, 'walk'): -1.0}

{(2, 'walk'): -6.983584610287968,
(3, 'tram'): -5.997740461664008,
(6, 'tram'): -3.504182096707191,
(12, 'walk'): -1.0}

src

def reinforcement_learning_magic_tram():
    from icecream import ic

    from ai.reinforcement.mdp_types import RandomPolicy
    from ai.reinforcement.mdp_example_tram import MDPMagicTram
    from ai.reinforcement.environment_from_mdp import EnvironmentFromMDP

    policy = RandomPolicy()
    mdp = MDPMagicTram(2, 13, fail_probability=0.2)
    environment = EnvironmentFromMDP(mdp)
    nb_episodes_per_state = 100_000
    start: State = 2
    ic(
        model_based_Monte_Carlo(
            environment, start, policy, nb_episodes_per_state=nb_episodes_per_state
        )
    )
    ic(
        model_free_Monte_Carlo(
            environment, start, policy, nb_episodes_per_state=nb_episodes_per_state
        )
    )
    ic(sarsa(environment, start, policy, nb_episodes_per_state=nb_episodes_per_state))
    opt_qualities, opt_policy_mapping = qlearning(
        environment, start, policy, nb_episodes_per_state=nb_episodes_per_state
    )
    ic(opt_qualities, opt_policy_mapping)

    opt_policy = PolicyFromMapping(opt_policy_mapping)
    ic(
        model_free_Monte_Carlo(
            environment, start, opt_policy, nb_episodes_per_state=nb_episodes_per_state
        )
    )
    ic(
        sarsa(
            environment, start, opt_policy, nb_episodes_per_state=nb_episodes_per_state
        )
    )

Swamp isle

Rearranging the Q-learning equation, one gets:

(1)\[\hat Q_{\textrm{opt}}^{(N+1)}(s, a) = \underbrace{\hat Q_{\textrm{opt}}^{(N)}(s, a)}_{\textrm{current}} + \underbrace{\alpha}_{\textrm{learning rate}} \left[ \underbrace{\left( r + \max_{a'} \gamma Q_{\textrm{opt}}^{(N)}(s', a') \right)}_{\textrm{target}} - \underbrace{\hat Q_{\textrm{opt}}^{(N)}(s, a)}_{\textrm{current}} \right]\]

This form is a classical form of iterative correction where the current value is updated proportionally to the difference between the better / target value and the current value. In this context, it is sometimes referred as temporal difference. In the original equation, a strict convex combination is equivalent to \(\alpha= \cfrac{1}{N+1}\). However, the next exercise will illustrate that this rapidly decreasing learning rate is not efficient.

Exercise (intermediate)

1. Implement and test the epsilon greedy algorithm (with defaulting to a random strategy when a state is not present in the strategy mapping).

2. Modify and test the qlearning algorithm in order to mutate continuously the exploration policy to the estimated optimum policy. Use the epsilon greedy policy for exploration, initialize its mapping with a non optimal policy. Comment on the results.

3. Improve the qlearning algorithm by using a fixed learning rate. Comment on the results.

Correction

1. The simplest experiment is to start with a failure probability of 0. This way the result should be identical to the state model and MDP exercises. Doin

2. Since the policy mapping is mutated along, there is no need to return a policy mapping at the end, the optimal policy is in the policy mapping of the policy. The counter (state, action) is no longer required because the learning rate is fixed.

   src

   class EpsilonGreedyPolicy(PolicyFromMapping):
       def __init__(self, policy_mapping: PolicyMapping, epsilon: float = 0.1):
           self.policy_mapping = dict(policy_mapping)
           self.epsilon = epsilon
   
       @override
       def action_from_state(self, state: State, actions: set[Action]) -> Action:
           actions_ = tuple(actions)
           opt_action: State | None = self.policy_mapping.get(state)
           if opt_action is None:
               return choice(actions_)  # random policy if no known optimal policy
               
           if len(actions)>1:
               non_opt_probability: float = self.epsilon / (len(actions) - 1)
           else:
               non_opt_probability: float = 0.0
               
           weights: list[float] = [
               non_opt_probability if action != opt_action else 1.0 - self.epsilon
               for action in actions_
           ] # for the case where len(actions)==1 => weights=[1-epsi]
           # this not a problem because choices normalizes the weights
           return choices(actions_, weights, k=1)[0]
   

   tests

   def test_epsilon_greedy():
       policy_mapping: PolicyMapping = {"in": "stay"}
       epsilon: float = 0.4
       nb_trials = 1000
       policy_epsilon_greedy = EpsilonGreedyPolicy(policy_mapping, epsilon=epsilon)
       are_actions_stay = [
           "stay" == policy_epsilon_greedy.action_from_state("in", {"stay", "quit"})
           for _ in range(nb_trials)
       ]
       assert pytest.approx(sum(are_actions_stay) / nb_trials, abs=0.05) == 1.0 - epsilon
   

   Even using a number of episodes per state very large (100_000 or even more), leads to very disappointing results: the qualities close to the destination are pretty well estimated, however, the qualities close to the starting position do not seen to converge. The problem with the decreasing learning rate is that the states close to the starting position are updated multiple times using very poor target values. The learning rate drops for these states, and later when target improves (from the destination), the corrections are too small. Using a random order for state traversal does not help to propagate the knowledge faster. With the fixed learning rate, the convergence is dramatically improved.

   {((0, 0), 'E'): -12.677577227122406,
   ((0, 0), 'S'): -12.667075392250679,
   ((0, 1), 'E'): -31.68226743698947,
   ((0, 1), 'S'): -31.66799828086022,
   ((0, 1), 'W'): -12.670359811413386,
   ((0, 2), 'E'): -40.9997934415334,
   ((0, 2), 'S'): -40.999834420809556,
   ((0, 2), 'W'): -31.67645420617531,
   ((0, 3), 'S'): -20.999992192948962,
   ((0, 3), 'W'): -50.681523379327565,
   ((1, 0), 'E'): -31.66673091925328,
   ((1, 0), 'N'): -12.66885326879473,
   ((1, 0), 'S'): -12.6661214455744,
   ((1, 1), 'E'): -40.9998103493078,
   ((1, 1), 'N'): -31.675785937100066,
   ((1, 1), 'S'): -31.66929228245872,
   ((1, 1), 'W'): -31.66414897696369,
   ((1, 2), 'E'): -21.000000000000004,
   ((1, 2), 'N'): -50.686172012334524,
   ((1, 2), 'S'): -39.999926341212834,
   ((1, 2), 'W'): -50.6678345317146,
   ((1, 3), 'N'): -21.999897582486756,
   ((1, 3), 'S'): -1.0,
   ((1, 3), 'W'): -21.999850806709116,
   ((2, 0), 'E'): -12.66813474888183,
   ((2, 0), 'N'): -12.667260215082385,
   ((2, 1), 'E'): -20.999777113079706,
   ((2, 1), 'N'): -31.66945779221507,
   ((2, 1), 'W'): -12.668045310032248,
   ((2, 2), 'E'): -19.999999999999993,
   ((2, 2), 'N'): -40.999794234146705,
   ((2, 2), 'W'): -31.667637590668363}
   
   {(0, 0): 'S',
   (0, 1): 'W',
   (0, 2): 'W',
   (0, 3): 'S',
   (1, 0): 'S',
   (1, 1): 'W',
   (1, 2): 'E',
   (1, 3): 'S',
   (2, 0): 'N',
   (2, 1): 'W',
   (2, 2): 'E'}
   
   {((0, 0), 'E'): -25.99999999999998,
   ((0, 0), 'S'): -23.999999999999982,
   ((0, 1), 'E'): -41.99999999999996,
   ((0, 1), 'S'): -41.99999999999996,
   ((0, 1), 'W'): -24.999999999999982,
   ((0, 2), 'E'): -40.99999999999997,
   ((0, 2), 'S'): -40.99999999999997,
   ((0, 2), 'W'): -44.99999999999999,
   ((0, 3), 'S'): -20.999999999999993,
   ((0, 3), 'W'): -60.99999999999998,
   ((1, 0), 'E'): -41.99999999999996,
   ((1, 0), 'N'): -24.999999999999982,
   ((1, 0), 'S'): -22.999999999999986,
   ((1, 1), 'E'): -40.99999999999997,
   ((1, 1), 'N'): -44.99999999999999,
   ((1, 1), 'S'): -40.99999999999997,
   ((1, 1), 'W'): -42.99999999999997,
   ((1, 2), 'E'): -20.999999999999993,
   ((1, 2), 'N'): -60.99999999999998,
   ((1, 2), 'S'): -39.999999999999986,
   ((1, 2), 'W'): -60.99999999999998,
   ((1, 3), 'N'): -21.99999999999999,
   ((1, 3), 'S'): -0.9999999999999998,
   ((1, 3), 'W'): -21.99999999999999,
   ((2, 0), 'E'): -21.99999999999999,
   ((2, 0), 'N'): -23.999999999999982,
   ((2, 1), 'E'): -20.999999999999993,
   ((2, 1), 'N'): -41.99999999999996,
   ((2, 1), 'W'): -22.999999999999986,
   ((2, 2), 'E'): -19.999999999999993,
   ((2, 2), 'N'): -40.99999999999997,
   ((2, 2), 'W'): -40.99999999999997}
   
   {(0, 0): 'S',
   (0, 1): 'W',
   (0, 2): 'S',
   (0, 3): 'S',
   (1, 0): 'S',
   (1, 1): 'S',
   (1, 2): 'E',
   (1, 3): 'S',
   (2, 0): 'E',
   (2, 1): 'E',
   (2, 2): 'E'}
   

3. src

   def qlearning_with_update(
       environment: Environment,
       state_: State,
       policy: PolicyFromMapping,  # change
       discount: float = 1.0,
       nb_episodes_per_state: int = 1_000,
       learning_rate: float = 0.2,  # change
   ) -> QualitiesMapping:  # change
       """contrary to the previous qlearning, this one receives the policy as input,
       and update the policy mapping along the way. For example, if the epsilon greedy
       polilcy is passed, obviously initially the optimal policy is not known (and defaulted to random).
       The optimal policy is thus gradually built and updated inside the policy.
       That is why the policy is not explictly returned."""
       states: set[State] = {state_}
       alpha = learning_rate
   
       opt_qualities: QualitiesMapping = defaultdict(float)
   
       # synchronize the optimum policy and exploration policy mappings
       opt_policy_mapping: PolicyMapping = policy.policy_mapping  # change
       states |= set(opt_policy_mapping)  # change
   
       # generate episode many times from all known states
       for _ in range(nb_episodes_per_state):
           for state in states.copy():
               episode: Generator[Frame] = environment.yield_frame_in_episode(
                   state, policy
               )
               if environment.is_end_state(state):
                   continue
   
               # analyze episode
               for frame in episode:
                   state, action, reward, next_state = frame
                   states.add(state)
                   states.add(next_state)
   
                   # obtain quality next state
                   if environment.is_end_state(next_state):
                       opt_next_quality: Reward = 0.0
                   else:
                       next_actions = environment.actions_from_state(next_state)
                       opt_next_quality = max(
                           [
                               opt_qualities[(next_state, next_action)]
                               for next_action in next_actions
                           ]
                       )
   
                   # update optimum qualities
                   opt_qualities[(state, action)] = (1.0 - alpha) * opt_qualities[
                       (state, action)
                   ] + alpha * (reward + discount * opt_next_quality)
   
                   # update optimum policy
                   actions = environment.actions_from_state(state)
                   _, opt_action = max(
                       [(opt_qualities[(state, action)], action) for action in actions]
                   )
                   opt_policy_mapping[state] = opt_action
   
       return opt_qualities