Repository: HumanCompatibleAI/overcooked_ai (language: python)

Issue:
Add `black` and `isort` to CI
Same as [this PR](https://github.com/HumanCompatibleAI/overcooked-demo/pull/38).


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/blackisort.yml b/.github/workflows/blackisort.yml
new file mode 100644
index 00000000..4115a9eb
--- /dev/null
+++ b/.github/workflows/blackisort.yml
@@ -0,0 +1,21 @@
+name: black isort
+
+on: [push, pull_request]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-python@v2
+      # https://black.readthedocs.io/en/stable/integrations/github_actions.html
+      - name: running Black
+        uses: psf/black@stable
+        with:
+          options: "--check --diff --line-length 79"
+      # https://pycqa.github.io/isort/docs/configuration/github_action.html
+      - name: running isort
+        uses: jamescurtin/isort-action@master
+        if: always()
+        with:
+          configuration: "--check-only --line-length 79 --profile black"
diff --git a/setup.py b/setup.py
index 4c7f4b33..5c27072c 100644
--- a/setup.py
+++ b/setup.py
@@ -1,36 +1,41 @@
 #!/usr/bin/env python
 
-from setuptools import setup, find_packages
+from setuptools import find_packages, setup
 
-with open("README.md", "r", encoding='UTF8') as fh:
+with open("README.md", "r", encoding="UTF8") as fh:
     long_description = fh.read()
 
-setup(name='overcooked_ai',
-      version='1.1.0',
-      description='Cooperative multi-agent environment based on Overcooked',
-      long_description=long_description,
-      long_description_content_type="text/markdown",
-      author='Micah Carroll',
-      author_email='mdc@berkeley.edu',
-      url='https://github.com/HumanCompatibleAI/overcooked_ai',
-      download_url='https://github.com/HumanCompatibleAI/overcooked_ai/archive/refs/tags/1.1.0.tar.gz',
-      packages=find_packages('src'),
-      keywords=['Overcooked', 'AI', 'Reinforcement Learning'],
-      package_dir={"": "src"},
-      package_data={
-        'overcooked_ai_py' : [
-          'data/layouts/*.layout', 'data/planners/*.py', 'data/human_data/*.pickle',
-          'data/graphics/*.png', 'data/graphics/*.json', 'data/fonts/*.ttf',
+setup(
+    name="overcooked_ai",
+    version="1.1.0",
+    description="Cooperative multi-agent environment based on Overcooked",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    author="Micah Carroll",
+    author_email="mdc@berkeley.edu",
+    url="https://github.com/HumanCompatibleAI/overcooked_ai",
+    download_url="https://github.com/HumanCompatibleAI/overcooked_ai/archive/refs/tags/1.1.0.tar.gz",
+    packages=find_packages("src"),
+    keywords=["Overcooked", "AI", "Reinforcement Learning"],
+    package_dir={"": "src"},
+    package_data={
+        "overcooked_ai_py": [
+            "data/layouts/*.layout",
+            "data/planners/*.py",
+            "data/human_data/*.pickle",
+            "data/graphics/*.png",
+            "data/graphics/*.json",
+            "data/fonts/*.ttf",
         ],
-      },
-      install_requires=[
-        'dill',
-        'numpy',
-        'scipy',
-        'tqdm',
-        'gym',
-        'ipython',
-        'pygame',
-        'ipywidgets'
-      ]
-    )
+    },
+    install_requires=[
+        "dill",
+        "numpy",
+        "scipy",
+        "tqdm",
+        "gym",
+        "ipython",
+        "pygame",
+        "ipywidgets",
+    ],
+)
diff --git a/src/overcooked_ai_py/__init__.py b/src/overcooked_ai_py/__init__.py
index bfa99163..42a19935 100644
--- a/src/overcooked_ai_py/__init__.py
+++ b/src/overcooked_ai_py/__init__.py
@@ -1,7 +1,6 @@
 from gym.envs.registration import register
 
 register(
-    id='Overcooked-v0',
-    entry_point='overcooked_ai_py.mdp.overcooked_env:Overcooked',
+    id="Overcooked-v0",
+    entry_point="overcooked_ai_py.mdp.overcooked_env:Overcooked",
 )
-
diff --git a/src/overcooked_ai_py/agents/agent.py b/src/overcooked_ai_py/agents/agent.py
index 281af257..425d9959 100644
--- a/src/overcooked_ai_py/agents/agent.py
+++ b/src/overcooked_ai_py/agents/agent.py
@@ -1,6 +1,11 @@
-import itertools, math, os, dill
-import numpy as np
+import itertools
+import math
+import os
 from collections import defaultdict
+
+import dill
+import numpy as np
+
 from overcooked_ai_py.mdp.actions import Action
 from overcooked_ai_py.mdp.overcooked_mdp import Recipe
 from overcooked_ai_py.utils import OvercookedException
@@ -8,7 +13,7 @@
 
 class Agent(object):
 
-    agent_file_name = 'agent.pickle'
+    agent_file_name = "agent.pickle"
 
     def __init__(self):
         self.reset()
@@ -27,7 +32,7 @@ def action(self, state):
     def actions(self, states, agent_indices):
         """
         A multi-state version of the action method. This enables for parallized
-        implementations that can potentially give speedups in action prediction. 
+        implementations that can potentially give speedups in action prediction.
 
         Args:
             states (list): list of OvercookedStates for which we want actions for
@@ -47,7 +52,11 @@ def a_probs_from_action(action):
     def check_action_probs(action_probs, tolerance=1e-4):
         """Check that action probabilities sum to ≈ 1.0"""
         probs_sum = sum(action_probs)
-        assert math.isclose(probs_sum, 1.0, rel_tol=tolerance), "Action probabilities {} should sum up to approximately 1 but sum up to {}".format(list(action_probs), probs_sum)
+        assert math.isclose(
+            probs_sum, 1.0, rel_tol=tolerance
+        ), "Action probabilities {} should sum up to approximately 1 but sum up to {}".format(
+            list(action_probs), probs_sum
+        )
 
     def set_agent_index(self, agent_index):
         self.agent_index = agent_index
@@ -65,11 +74,13 @@ def reset(self):
 
     def save(self, path):
         if os.path.isfile(path):
-            raise IOError("Must specify a path to directory! Got: {}".format(path))
+            raise IOError(
+                "Must specify a path to directory! Got: {}".format(path)
+            )
         if not os.path.exists(path):
             os.makedirs(path)
         pickle_path = os.path.join(path, self.agent_file_name)
-        with open(pickle_path, 'wb') as f:
+        with open(pickle_path, "wb") as f:
             dill.dump(self, f)
         return path
 
@@ -78,19 +89,19 @@ def load(cls, path):
         if os.path.isdir(path):
             path = os.path.join(path, cls.agent_file_name)
         try:
-            with open(path, 'rb') as f:
+            with open(path, "rb") as f:
                 obj = dill.load(f)
             return obj
         except OvercookedException:
             Recipe.configure({})
-            with open(path, 'rb') as f:
+            with open(path, "rb") as f:
                 obj = dill.load(f)
             return obj
 
 
 class AgentGroup(object):
     """
-    AgentGroup is a group of N agents used to sample 
+    AgentGroup is a group of N agents used to sample
     joint actions in the context of an OvercookedEnv instance.
     """
 
@@ -99,8 +110,12 @@ def __init__(self, *agents, allow_duplicate_agents=False):
         self.n = len(self.agents)
         self.reset()
 
-        if not all(a0 is not a1 for a0, a1 in itertools.combinations(agents, 2)):
-            assert allow_duplicate_agents, "All agents should be separate instances, unless allow_duplicate_agents is set to true"
+        if not all(
+            a0 is not a1 for a0, a1 in itertools.combinations(agents, 2)
+        ):
+            assert (
+                allow_duplicate_agents
+            ), "All agents should be separate instances, unless allow_duplicate_agents is set to true"
 
     def joint_action(self, state):
         actions_and_probs_n = tuple(a.action(state) for a in self.agents)
@@ -129,8 +144,10 @@ class AgentPair(AgentGroup):
     for both fields can lead to problems if the agents have state / history)
     """
 
-    def __init__(self, *agents, allow_duplicate_agents=False): 
-        super().__init__(*agents, allow_duplicate_agents=allow_duplicate_agents)
+    def __init__(self, *agents, allow_duplicate_agents=False):
+        super().__init__(
+            *agents, allow_duplicate_agents=allow_duplicate_agents
+        )
         assert self.n == 2
         self.a0, self.a1 = self.agents
 
@@ -147,6 +164,7 @@ def joint_action(self, state):
         else:
             return super().joint_action(state)
 
+
 class NNPolicy(object):
     """
     This is a common format for NN-based policies. Once one has wrangled the intended trained neural net
@@ -173,14 +191,14 @@ class AgentFromPolicy(Agent):
     """
     This is a useful Agent class backbone from which to subclass from NN-based agents.
     """
-    
+
     def __init__(self, policy):
         """
         Takes as input an NN Policy instance
         """
         self.policy = policy
         self.reset()
-        
+
     def action(self, state):
         return self.actions([state], [self.agent_index])[0]
 
@@ -189,13 +207,15 @@ def actions(self, states, agent_indices):
         actions_and_infos_n = []
         for action_probs in action_probs_n:
             action = Action.sample(action_probs)
-            actions_and_infos_n.append((action, {"action_probs": action_probs}))
+            actions_and_infos_n.append(
+                (action, {"action_probs": action_probs})
+            )
         return actions_and_infos_n
 
     def set_mdp(self, mdp):
         super().set_mdp(mdp)
         self.policy.mdp = mdp
-    
+
     def reset(self):
         super(AgentFromPolicy, self).reset()
         self.policy.mdp = None
@@ -207,17 +227,21 @@ class RandomAgent(Agent):
     NOTE: Does not perform interact actions, unless specified
     """
 
-    def __init__(self, sim_threads=None, all_actions=False, custom_wait_prob=None):
+    def __init__(
+        self, sim_threads=None, all_actions=False, custom_wait_prob=None
+    ):
         self.sim_threads = sim_threads
         self.all_actions = all_actions
         self.custom_wait_prob = custom_wait_prob
-    
+
     def action(self, state):
         action_probs = np.zeros(Action.NUM_ACTIONS)
         legal_actions = list(Action.MOTION_ACTIONS)
         if self.all_actions:
             legal_actions = Action.ALL_ACTIONS
-        legal_actions_indices = np.array([Action.ACTION_TO_INDEX[motion_a] for motion_a in legal_actions])
+        legal_actions_indices = np.array(
+            [Action.ACTION_TO_INDEX[motion_a] for motion_a in legal_actions]
+        )
         action_probs[legal_actions_indices] = 1 / len(legal_actions_indices)
 
         if self.custom_wait_prob is not None:
@@ -225,7 +249,9 @@ def action(self, state):
             if np.random.random() < self.custom_wait_prob:
                 return stay, {"action_probs": Agent.a_probs_from_action(stay)}
             else:
-                action_probs = Action.remove_indices_and_renormalize(action_probs, [Action.ACTION_TO_INDEX[stay]])
+                action_probs = Action.remove_indices_and_renormalize(
+                    action_probs, [Action.ACTION_TO_INDEX[stay]]
+                )
 
         return Action.sample(action_probs), {"action_probs": action_probs}
 
@@ -237,10 +263,9 @@ def direct_action(self, obs):
 
 
 class StayAgent(Agent):
-
     def __init__(self, sim_threads=None):
         self.sim_threads = sim_threads
-    
+
     def action(self, state):
         a = Action.STAY
         return a, {}
@@ -258,7 +283,7 @@ class FixedPlanAgent(Agent):
     def __init__(self, plan):
         self.plan = plan
         self.i = 0
-    
+
     def action(self, state):
         if self.i >= len(self.plan):
             return Action.STAY, {}
@@ -281,14 +306,23 @@ class GreedyHumanModel(Agent):
     Will work only in environments where the only order is 3 onion soup.
     """
 
-    def __init__(self, mlam, hl_boltzmann_rational=False, ll_boltzmann_rational=False, hl_temp=1, ll_temp=1,
-                 auto_unstuck=True):
+    def __init__(
+        self,
+        mlam,
+        hl_boltzmann_rational=False,
+        ll_boltzmann_rational=False,
+        hl_temp=1,
+        ll_temp=1,
+        auto_unstuck=True,
+    ):
         self.mlam = mlam
         self.mdp = self.mlam.mdp
 
         # Bool for perfect rationality vs Boltzmann rationality for high level and low level action selection
         self.hl_boltzmann_rational = hl_boltzmann_rational  # For choices among high level goals of same type
-        self.ll_boltzmann_rational = ll_boltzmann_rational  # For choices about low level motion
+        self.ll_boltzmann_rational = (
+            ll_boltzmann_rational  # For choices about low level motion
+        )
 
         # Coefficient for Boltzmann rationality for high level action selection
         self.hl_temperature = hl_temp
@@ -318,32 +352,53 @@ def action(self, state):
         # level action we want to perform, select the one with lowest cost
         start_pos_and_or = state.players_pos_and_or[self.agent_index]
 
-        chosen_goal, chosen_action, action_probs = self.choose_motion_goal(start_pos_and_or, possible_motion_goals)
+        chosen_goal, chosen_action, action_probs = self.choose_motion_goal(
+            start_pos_and_or, possible_motion_goals
+        )
 
-        if self.ll_boltzmann_rational and chosen_goal[0] == start_pos_and_or[0]:
-            chosen_action, action_probs = self.boltzmann_rational_ll_action(start_pos_and_or, chosen_goal)
+        if (
+            self.ll_boltzmann_rational
+            and chosen_goal[0] == start_pos_and_or[0]
+        ):
+            chosen_action, action_probs = self.boltzmann_rational_ll_action(
+                start_pos_and_or, chosen_goal
+            )
 
         if self.auto_unstuck:
             # HACK: if two agents get stuck, select an action at random that would
             # change the player positions if the other player were not to move
-            if self.prev_state is not None and state.players_pos_and_or == self.prev_state.players_pos_and_or:
+            if (
+                self.prev_state is not None
+                and state.players_pos_and_or
+                == self.prev_state.players_pos_and_or
+            ):
                 if self.agent_index == 0:
-                    joint_actions = list(itertools.product(Action.ALL_ACTIONS, [Action.STAY]))
+                    joint_actions = list(
+                        itertools.product(Action.ALL_ACTIONS, [Action.STAY])
+                    )
                 elif self.agent_index == 1:
-                    joint_actions = list(itertools.product([Action.STAY], Action.ALL_ACTIONS))
+                    joint_actions = list(
+                        itertools.product([Action.STAY], Action.ALL_ACTIONS)
+                    )
                 else:
                     raise ValueError("Player index not recognized")
 
                 unblocking_joint_actions = []
                 for j_a in joint_actions:
-                    new_state, _ = self.mlam.mdp.get_state_transition(state, j_a)
-                    if new_state.player_positions != self.prev_state.player_positions:
+                    new_state, _ = self.mlam.mdp.get_state_transition(
+                        state, j_a
+                    )
+                    if (
+                        new_state.player_positions
+                        != self.prev_state.player_positions
+                    ):
                         unblocking_joint_actions.append(j_a)
                 # Getting stuck became a possiblity simply because the nature of a layout (having a dip in the middle)
                 if len(unblocking_joint_actions) == 0:
                     unblocking_joint_actions.append([Action.STAY, Action.STAY])
-                chosen_action = unblocking_joint_actions[np.random.choice(len(unblocking_joint_actions))][
-                    self.agent_index]
+                chosen_action = unblocking_joint_actions[
+                    np.random.choice(len(unblocking_joint_actions))
+                ][self.agent_index]
                 action_probs = self.a_probs_from_action(chosen_action)
 
             # NOTE: Assumes that calls to the action method are sequential
@@ -357,20 +412,32 @@ def choose_motion_goal(self, start_pos_and_or, motion_goals):
         or rationally), and returns the plan and the corresponding first action on that plan.
         """
         if self.hl_boltzmann_rational:
-            possible_plans = [self.mlam.motion_planner.get_plan(start_pos_and_or, goal) for goal in motion_goals]
+            possible_plans = [
+                self.mlam.motion_planner.get_plan(start_pos_and_or, goal)
+                for goal in motion_goals
+            ]
             plan_costs = [plan[2] for plan in possible_plans]
-            goal_idx, action_probs = self.get_boltzmann_rational_action_idx(plan_costs, self.hl_temperature)
+            goal_idx, action_probs = self.get_boltzmann_rational_action_idx(
+                plan_costs, self.hl_temperature
+            )
             chosen_goal = motion_goals[goal_idx]
             chosen_goal_action = possible_plans[goal_idx][0][0]
         else:
-            chosen_goal, chosen_goal_action = self.get_lowest_cost_action_and_goal(start_pos_and_or, motion_goals)
+            (
+                chosen_goal,
+                chosen_goal_action,
+            ) = self.get_lowest_cost_action_and_goal(
+                start_pos_and_or, motion_goals
+            )
             action_probs = self.a_probs_from_action(chosen_goal_action)
         return chosen_goal, chosen_goal_action, action_probs
 
     def get_boltzmann_rational_action_idx(self, costs, temperature):
         """Chooses index based on softmax probabilities obtained from cost array"""
         costs = np.array(costs)
-        softmax_probs = np.exp(-costs * temperature) / np.sum(np.exp(-costs * temperature))
+        softmax_probs = np.exp(-costs * temperature) / np.sum(
+            np.exp(-costs * temperature)
+        )
         action_idx = np.random.choice(len(costs), p=softmax_probs)
         return action_idx, softmax_probs
 
@@ -382,14 +449,18 @@ def get_lowest_cost_action_and_goal(self, start_pos_and_or, motion_goals):
         min_cost = np.Inf
         best_action, best_goal = None, None
         for goal in motion_goals:
-            action_plan, _, plan_cost = self.mlam.motion_planner.get_plan(start_pos_and_or, goal)
+            action_plan, _, plan_cost = self.mlam.motion_planner.get_plan(
+                start_pos_and_or, goal
+            )
             if plan_cost < min_cost:
                 best_action = action_plan[0]
                 min_cost = plan_cost
                 best_goal = goal
         return best_goal, best_action
 
-    def boltzmann_rational_ll_action(self, start_pos_and_or, goal, inverted_costs=False):
+    def boltzmann_rational_ll_action(
+        self, start_pos_and_or, goal, inverted_costs=False
+    ):
         """
         Computes the plan cost to reach the goal after taking each possible low level action.
         Selects a low level action boltzmann rationally based on the one-step-ahead plan costs.
@@ -401,11 +472,15 @@ def boltzmann_rational_ll_action(self, start_pos_and_or, goal, inverted_costs=Fa
         for action in Action.ALL_ACTIONS:
             pos, orient = start_pos_and_or
             new_pos_and_or = self.mdp._move_if_direction(pos, orient, action)
-            _, _, plan_cost = self.mlam.motion_planner.get_plan(new_pos_and_or, goal)
+            _, _, plan_cost = self.mlam.motion_planner.get_plan(
+                new_pos_and_or, goal
+            )
             sign = (-1) ** int(inverted_costs)
             future_costs.append(sign * plan_cost)
 
-        action_idx, action_probs = self.get_boltzmann_rational_action_idx(future_costs, self.ll_temperature)
+        action_idx, action_probs = self.get_boltzmann_rational_action_idx(
+            future_costs, self.ll_temperature
+        )
         return Action.ALL_ACTIONS[action_idx], action_probs
 
     def ml_action(self, state):
@@ -424,31 +499,44 @@ def ml_action(self, state):
         other_player = state.players[1 - self.agent_index]
         am = self.mlam
 
-        counter_objects = self.mlam.mdp.get_counter_objects_dict(state, list(self.mlam.mdp.terrain_pos_dict['X']))
+        counter_objects = self.mlam.mdp.get_counter_objects_dict(
+            state, list(self.mlam.mdp.terrain_pos_dict["X"])
+        )
         pot_states_dict = self.mlam.mdp.get_pot_states(state)
 
-
         if not player.has_object():
-            ready_soups = pot_states_dict['ready']
-            cooking_soups = pot_states_dict['cooking']
+            ready_soups = pot_states_dict["ready"]
+            cooking_soups = pot_states_dict["cooking"]
 
             soup_nearly_ready = len(ready_soups) > 0 or len(cooking_soups) > 0
-            other_has_dish = other_player.has_object() and other_player.get_object().name == 'dish'
+            other_has_dish = (
+                other_player.has_object()
+                and other_player.get_object().name == "dish"
+            )
 
             if soup_nearly_ready and not other_has_dish:
                 motion_goals = am.pickup_dish_actions(counter_objects)
             else:
-                assert len(state.all_orders) == 1 and list(state.all_orders[0].ingredients) == ["onion", "onion", "onion"], \
-                    "The current mid level action manager only support 3-onion-soup order, but got orders" \
+                assert len(state.all_orders) == 1 and list(
+                    state.all_orders[0].ingredients
+                ) == ["onion", "onion", "onion"], (
+                    "The current mid level action manager only support 3-onion-soup order, but got orders"
                     + str(state.all_orders)
+                )
                 next_order = list(state.all_orders)[0]
-                soups_ready_to_cook_key = '{}_items'.format(len(next_order.ingredients))
+                soups_ready_to_cook_key = "{}_items".format(
+                    len(next_order.ingredients)
+                )
                 soups_ready_to_cook = pot_states_dict[soups_ready_to_cook_key]
                 if soups_ready_to_cook:
                     only_pot_states_ready_to_cook = defaultdict(list)
-                    only_pot_states_ready_to_cook[soups_ready_to_cook_key] = soups_ready_to_cook
+                    only_pot_states_ready_to_cook[
+                        soups_ready_to_cook_key
+                    ] = soups_ready_to_cook
                     # we want to cook only soups that has same len as order
-                    motion_goals = am.start_cooking_actions(only_pot_states_ready_to_cook)
+                    motion_goals = am.start_cooking_actions(
+                        only_pot_states_ready_to_cook
+                    )
                 else:
                     motion_goals = am.pickup_onion_actions(counter_objects)
                 # it does not make sense to have tomato logic when the only possible order is 3 onion soup (see assertion above)
@@ -459,38 +547,51 @@ def ml_action(self, state):
                 # else:
                 #     motion_goals = am.pickup_onion_actions(counter_objects) + am.pickup_tomato_actions(counter_objects)
 
-
         else:
             player_obj = player.get_object()
 
-            if player_obj.name == 'onion':
+            if player_obj.name == "onion":
                 motion_goals = am.put_onion_in_pot_actions(pot_states_dict)
 
-            elif player_obj.name == 'tomato':
+            elif player_obj.name == "tomato":
                 motion_goals = am.put_tomato_in_pot_actions(pot_states_dict)
 
-            elif player_obj.name == 'dish':
-                motion_goals = am.pickup_soup_with_dish_actions(pot_states_dict, only_nearly_ready=True)
+            elif player_obj.name == "dish":
+                motion_goals = am.pickup_soup_with_dish_actions(
+                    pot_states_dict, only_nearly_ready=True
+                )
 
-            elif player_obj.name == 'soup':
+            elif player_obj.name == "soup":
                 motion_goals = am.deliver_soup_actions()
 
             else:
                 raise ValueError()
 
-        motion_goals = [mg for mg in motion_goals if self.mlam.motion_planner.is_valid_motion_start_goal_pair(player.pos_and_or, mg)]
+        motion_goals = [
+            mg
+            for mg in motion_goals
+            if self.mlam.motion_planner.is_valid_motion_start_goal_pair(
+                player.pos_and_or, mg
+            )
+        ]
 
         if len(motion_goals) == 0:
             motion_goals = am.go_to_closest_feature_actions(player)
-            motion_goals = [mg for mg in motion_goals if self.mlam.motion_planner.is_valid_motion_start_goal_pair(player.pos_and_or, mg)]
+            motion_goals = [
+                mg
+                for mg in motion_goals
+                if self.mlam.motion_planner.is_valid_motion_start_goal_pair(
+                    player.pos_and_or, mg
+                )
+            ]
             assert len(motion_goals) != 0
 
         return motion_goals
 
 
 class SampleAgent(Agent):
-    """ Agent that samples action using the average action_probs across multiple agents
-    """
+    """Agent that samples action using the average action_probs across multiple agents"""
+
     def __init__(self, agents):
         self.agents = agents
 
@@ -498,7 +599,7 @@ def action(self, state):
         action_probs = np.zeros(Action.NUM_ACTIONS)
         for agent in self.agents:
             action_probs += agent.action(state)[1]["action_probs"]
-        action_probs = action_probs/len(self.agents)
+        action_probs = action_probs / len(self.agents)
         return Action.sample(action_probs), {"action_probs": action_probs}
 
 
diff --git a/src/overcooked_ai_py/agents/benchmarking.py b/src/overcooked_ai_py/agents/benchmarking.py
index 4e3fa98f..5c40903e 100644
--- a/src/overcooked_ai_py/agents/benchmarking.py
+++ b/src/overcooked_ai_py/agents/benchmarking.py
@@ -1,14 +1,32 @@
 import copy
+
 import numpy as np
 
-from overcooked_ai_py.utils import save_pickle, load_pickle, cumulative_rewards_from_rew_list, save_as_json, \
-    load_from_json, merge_dictionaries, rm_idx_from_dict, take_indexes_from_dict, is_iterable
-from overcooked_ai_py.planning.planners import NO_COUNTERS_PARAMS
-from overcooked_ai_py.agents.agent import AgentPair, RandomAgent, GreedyHumanModel
-from overcooked_ai_py.mdp.overcooked_mdp import OvercookedGridworld, Action, OvercookedState
-from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
+from overcooked_ai_py.agents.agent import (
+    AgentPair,
+    GreedyHumanModel,
+    RandomAgent,
+)
 from overcooked_ai_py.mdp.layout_generator import LayoutGenerator
+from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
+from overcooked_ai_py.mdp.overcooked_mdp import (
+    Action,
+    OvercookedGridworld,
+    OvercookedState,
+)
 from overcooked_ai_py.mdp.overcooked_trajectory import DEFAULT_TRAJ_KEYS
+from overcooked_ai_py.planning.planners import NO_COUNTERS_PARAMS
+from overcooked_ai_py.utils import (
+    cumulative_rewards_from_rew_list,
+    is_iterable,
+    load_from_json,
+    load_pickle,
+    merge_dictionaries,
+    rm_idx_from_dict,
+    save_as_json,
+    save_pickle,
+    take_indexes_from_dict,
+)
 
 
 class AgentEvaluator(object):
@@ -21,7 +39,14 @@ class AgentEvaluator(object):
     pickleable. We should think about possible improvements/what makes most sense to do here.
     """
 
-    def __init__(self, env_params, mdp_fn, force_compute=False, mlam_params=NO_COUNTERS_PARAMS, debug=False):
+    def __init__(
+        self,
+        env_params,
+        mdp_fn,
+        force_compute=False,
+        mlam_params=NO_COUNTERS_PARAMS,
+        debug=False,
+    ):
         """
         env_params (dict): params for creation of an OvercookedEnv
         mdp_fn (callable function): a function that can be used to create mdp
@@ -29,15 +54,24 @@ def __init__(self, env_params, mdp_fn, force_compute=False, mlam_params=NO_COUNT
         mlam_params (dict): the parameters for mlam, the MediumLevelActionManager
         debug (bool): whether to display debugging information on init
         """
-        assert callable(mdp_fn), "mdp generating function must be a callable function"
+        assert callable(
+            mdp_fn
+        ), "mdp generating function must be a callable function"
         env_params["mlam_params"] = mlam_params
         self.mdp_fn = mdp_fn
         self.env = OvercookedEnv(self.mdp_fn, **env_params)
         self.force_compute = force_compute
 
     @staticmethod
-    def from_mdp_params_infinite(mdp_params, env_params, outer_shape=None, mdp_params_schedule_fn=None, force_compute=False,
-                        mlam_params=NO_COUNTERS_PARAMS, debug=False):
+    def from_mdp_params_infinite(
+        mdp_params,
+        env_params,
+        outer_shape=None,
+        mdp_params_schedule_fn=None,
+        force_compute=False,
+        mlam_params=NO_COUNTERS_PARAMS,
+        debug=False,
+    ):
         """
         mdp_params (dict): params for creation of an OvercookedGridworld instance through the `from_layout_name` method
         outer_shape: the outer shape of environment
@@ -46,15 +80,29 @@ def from_mdp_params_infinite(mdp_params, env_params, outer_shape=None, mdp_param
 
         Infinitely generate mdp using the naive mdp_fn
         """
-        assert outer_shape is not None, "outer_shape needs to be defined for variable mdp"
-        assert "num_mdp" in env_params and np.isinf(env_params["num_mdp"]), \
-            "num_mdp needs to be specified and infinite"
-        mdp_fn_naive = LayoutGenerator.mdp_gen_fn_from_dict(mdp_params, outer_shape, mdp_params_schedule_fn)
-        return AgentEvaluator(env_params, mdp_fn_naive, force_compute, mlam_params, debug)
+        assert (
+            outer_shape is not None
+        ), "outer_shape needs to be defined for variable mdp"
+        assert "num_mdp" in env_params and np.isinf(
+            env_params["num_mdp"]
+        ), "num_mdp needs to be specified and infinite"
+        mdp_fn_naive = LayoutGenerator.mdp_gen_fn_from_dict(
+            mdp_params, outer_shape, mdp_params_schedule_fn
+        )
+        return AgentEvaluator(
+            env_params, mdp_fn_naive, force_compute, mlam_params, debug
+        )
 
     @staticmethod
-    def from_mdp_params_finite(mdp_params, env_params, outer_shape=None, mdp_params_schedule_fn=None, force_compute=False,
-                        mlam_params=NO_COUNTERS_PARAMS, debug=False):
+    def from_mdp_params_finite(
+        mdp_params,
+        env_params,
+        outer_shape=None,
+        mdp_params_schedule_fn=None,
+        force_compute=False,
+        mlam_params=NO_COUNTERS_PARAMS,
+        debug=False,
+    ):
         """
         mdp_params (dict): params for creation of an OvercookedGridworld instance through the `from_layout_name` method
         outer_shape: the outer shape of environment
@@ -64,65 +112,135 @@ def from_mdp_params_finite(mdp_params, env_params, outer_shape=None, mdp_params_
         Generate a finite list of mdp (mdp_lst) using the naive mdp_fn, and then use the from_mdp_lst to generate
         the AgentEvaluator
         """
-        assert outer_shape is not None, "outer_shape needs to be defined for variable mdp"
-        assert "num_mdp" in env_params and not np.isinf(env_params["num_mdp"]), \
-            "num_mdp needs to be specified and finite"
-        mdp_fn_naive = LayoutGenerator.mdp_gen_fn_from_dict(mdp_params, outer_shape, mdp_params_schedule_fn)
+        assert (
+            outer_shape is not None
+        ), "outer_shape needs to be defined for variable mdp"
+        assert "num_mdp" in env_params and not np.isinf(
+            env_params["num_mdp"]
+        ), "num_mdp needs to be specified and finite"
+        mdp_fn_naive = LayoutGenerator.mdp_gen_fn_from_dict(
+            mdp_params, outer_shape, mdp_params_schedule_fn
+        )
         # finite mdp, random choice
-        num_mdp = env_params['num_mdp']
-        assert type(num_mdp) == int and num_mdp > 0, "invalid number of mdp: " + str(num_mdp)
+        num_mdp = env_params["num_mdp"]
+        assert (
+            type(num_mdp) == int and num_mdp > 0
+        ), "invalid number of mdp: " + str(num_mdp)
         mdp_lst = [mdp_fn_naive() for _ in range(num_mdp)]
-        return AgentEvaluator.from_mdp_lst(mdp_lst=mdp_lst, env_params=env_params,
-                                           force_compute=force_compute, mlam_params=mlam_params, debug=debug)
+        return AgentEvaluator.from_mdp_lst(
+            mdp_lst=mdp_lst,
+            env_params=env_params,
+            force_compute=force_compute,
+            mlam_params=mlam_params,
+            debug=debug,
+        )
 
     @staticmethod
-    def from_mdp(mdp, env_params, force_compute=False, mlam_params=NO_COUNTERS_PARAMS, debug=False):
+    def from_mdp(
+        mdp,
+        env_params,
+        force_compute=False,
+        mlam_params=NO_COUNTERS_PARAMS,
+        debug=False,
+    ):
         """
         mdp (OvercookedGridworld): the mdp that we want the AgentEvaluator to always generate
         Information for the rest of params please refer to the __init__ method above
         """
-        assert type(mdp) == OvercookedGridworld, "mdp must be a OvercookedGridworld object"
+        assert (
+            type(mdp) == OvercookedGridworld
+        ), "mdp must be a OvercookedGridworld object"
         mdp_fn = lambda _ignored: mdp
-        return AgentEvaluator(env_params, mdp_fn, force_compute, mlam_params, debug)
+        return AgentEvaluator(
+            env_params, mdp_fn, force_compute, mlam_params, debug
+        )
 
     @staticmethod
-    def from_layout_name(mdp_params, env_params, force_compute=False, mlam_params=NO_COUNTERS_PARAMS, debug=False):
+    def from_layout_name(
+        mdp_params,
+        env_params,
+        force_compute=False,
+        mlam_params=NO_COUNTERS_PARAMS,
+        debug=False,
+    ):
         """
         mdp_params (dict): params for creation of an OvercookedGridworld instance through the `from_layout_name` method
         Information for the rest of params please refer to the __init__ method above
         """
         assert type(mdp_params) is dict and "layout_name" in mdp_params
         mdp = OvercookedGridworld.from_layout_name(**mdp_params)
-        return AgentEvaluator.from_mdp(mdp, env_params, force_compute, mlam_params, debug)
+        return AgentEvaluator.from_mdp(
+            mdp, env_params, force_compute, mlam_params, debug
+        )
 
     @staticmethod
-    def from_mdp_lst(mdp_lst, env_params, sampling_freq=None, force_compute=False, mlam_params=NO_COUNTERS_PARAMS, debug=False):
+    def from_mdp_lst(
+        mdp_lst,
+        env_params,
+        sampling_freq=None,
+        force_compute=False,
+        mlam_params=NO_COUNTERS_PARAMS,
+        debug=False,
+    ):
         """
         mdp_lst (list): a list of mdp (OvercookedGridworld) we would like to
         sampling_freq (list): a list of number that signify the sampling frequency of each mdp in the mdp_lst
         Information for the rest of params please refer to the __init__ method above
         """
         assert is_iterable(mdp_lst), "mdp_lst must be a list"
-        assert all([type(mdp) == OvercookedGridworld for mdp in mdp_lst]), "some mdps are not OvercookedGridworld objects"
+        assert all(
+            [type(mdp) == OvercookedGridworld for mdp in mdp_lst]
+        ), "some mdps are not OvercookedGridworld objects"
 
         if sampling_freq is None:
-            sampling_freq = np.ones(len(mdp_lst)) /len(mdp_lst)
+            sampling_freq = np.ones(len(mdp_lst)) / len(mdp_lst)
 
         mdp_fn = lambda _ignored: np.random.choice(mdp_lst, p=sampling_freq)
-        return AgentEvaluator(env_params, mdp_fn, force_compute, mlam_params, debug)
+        return AgentEvaluator(
+            env_params, mdp_fn, force_compute, mlam_params, debug
+        )
 
-    def evaluate_random_pair(self, num_games=1, all_actions=True, display=False, native_eval=False):
-        agent_pair = AgentPair(RandomAgent(all_actions=all_actions), RandomAgent(all_actions=all_actions))
-        return self.evaluate_agent_pair(agent_pair, num_games=num_games, display=display, native_eval=native_eval)
+    def evaluate_random_pair(
+        self, num_games=1, all_actions=True, display=False, native_eval=False
+    ):
+        agent_pair = AgentPair(
+            RandomAgent(all_actions=all_actions),
+            RandomAgent(all_actions=all_actions),
+        )
+        return self.evaluate_agent_pair(
+            agent_pair,
+            num_games=num_games,
+            display=display,
+            native_eval=native_eval,
+        )
 
-    def evaluate_human_model_pair(self, num_games=1, display=False, native_eval=False):
+    def evaluate_human_model_pair(
+        self, num_games=1, display=False, native_eval=False
+    ):
         a0 = GreedyHumanModel(self.env.mlam)
         a1 = GreedyHumanModel(self.env.mlam)
         agent_pair = AgentPair(a0, a1)
-        return self.evaluate_agent_pair(agent_pair, num_games=num_games, display=display, native_eval=native_eval)
+        return self.evaluate_agent_pair(
+            agent_pair,
+            num_games=num_games,
+            display=display,
+            native_eval=native_eval,
+        )
 
-    def evaluate_agent_pair(self, agent_pair, num_games, game_length=None, start_state_fn=None, metadata_fn=None, metadata_info_fn=None, display=False, dir=None,
-                            display_phi=False, info=True, native_eval=False):
+    def evaluate_agent_pair(
+        self,
+        agent_pair,
+        num_games,
+        game_length=None,
+        start_state_fn=None,
+        metadata_fn=None,
+        metadata_info_fn=None,
+        display=False,
+        dir=None,
+        display_phi=False,
+        info=True,
+        native_eval=False,
+    ):
         # this index has to be 0 because the Agent_Evaluator only has 1 env initiated
         # if you would like to evaluate on a different env using rllib, please modifiy
         # rllib/ -> rllib.py -> get_rllib_eval_function -> _evaluate
@@ -131,24 +249,76 @@ def evaluate_agent_pair(self, agent_pair, num_games, game_length=None, start_sta
         # this is particulally helpful with variable MDP, where we want to make sure
         # the mdp used in evaluation is the same as the native self.env.mdp
         if native_eval:
-            return self.env.get_rollouts(agent_pair, num_games=num_games, display=display, dir=dir, display_phi=display_phi,
-                                         info=info, metadata_fn=metadata_fn, metadata_info_fn=metadata_info_fn)
+            return self.env.get_rollouts(
+                agent_pair,
+                num_games=num_games,
+                display=display,
+                dir=dir,
+                display_phi=display_phi,
+                info=info,
+                metadata_fn=metadata_fn,
+                metadata_info_fn=metadata_info_fn,
+            )
         else:
             horizon_env = self.env.copy()
-            horizon_env.horizon = self.env.horizon if game_length is None else game_length
-            horizon_env.start_state_fn = self.env.start_state_fn if start_state_fn is None else start_state_fn
+            horizon_env.horizon = (
+                self.env.horizon if game_length is None else game_length
+            )
+            horizon_env.start_state_fn = (
+                self.env.start_state_fn
+                if start_state_fn is None
+                else start_state_fn
+            )
             horizon_env.reset()
-            return horizon_env.get_rollouts(agent_pair, num_games=num_games, display=display, dir=dir, display_phi=display_phi,
-                                            info=info, metadata_fn=metadata_fn, metadata_info_fn=metadata_info_fn)
+            return horizon_env.get_rollouts(
+                agent_pair,
+                num_games=num_games,
+                display=display,
+                dir=dir,
+                display_phi=display_phi,
+                info=info,
+                metadata_fn=metadata_fn,
+                metadata_info_fn=metadata_info_fn,
+            )
 
-    def get_agent_pair_trajs(self, a0, a1=None, num_games=100, game_length=None, start_state_fn=None, display=False, info=True):
+    def get_agent_pair_trajs(
+        self,
+        a0,
+        a1=None,
+        num_games=100,
+        game_length=None,
+        start_state_fn=None,
+        display=False,
+        info=True,
+    ):
         """Evaluate agent pair on both indices, and return trajectories by index"""
         if a1 is None:
             ap = AgentPair(a0, a0, allow_duplicate_agents=True)
-            trajs_0 = trajs_1 = self.evaluate_agent_pair(ap, num_games=num_games, game_length=game_length, start_state_fn=start_state_fn, display=display, info=info)
+            trajs_0 = trajs_1 = self.evaluate_agent_pair(
+                ap,
+                num_games=num_games,
+                game_length=game_length,
+                start_state_fn=start_state_fn,
+                display=display,
+                info=info,
+            )
         else:
-            trajs_0 = self.evaluate_agent_pair(AgentPair(a0, a1), num_games=num_games, game_length=game_length, start_state_fn=start_state_fn, display=display, info=info)
-            trajs_1 = self.evaluate_agent_pair(AgentPair(a1, a0), num_games=num_games, game_length=game_length, start_state_fn=start_state_fn, display=display, info=info)
+            trajs_0 = self.evaluate_agent_pair(
+                AgentPair(a0, a1),
+                num_games=num_games,
+                game_length=game_length,
+                start_state_fn=start_state_fn,
+                display=display,
+                info=info,
+            )
+            trajs_1 = self.evaluate_agent_pair(
+                AgentPair(a1, a0),
+                num_games=num_games,
+                game_length=game_length,
+                start_state_fn=start_state_fn,
+                display=display,
+                info=info,
+            )
         return trajs_0, trajs_1
 
     @staticmethod
@@ -167,15 +337,25 @@ def check_trajectories(trajectories, from_json=False, **kwargs):
     @staticmethod
     def _check_standard_traj_keys(traj_keys_set):
         default_traj_keys = DEFAULT_TRAJ_KEYS
-        assert traj_keys_set == set(default_traj_keys), "Keys of traj dict did not match standard form.\nMissing keys: {}\nAdditional keys: {}".format(
-            [k for k in default_traj_keys if k not in traj_keys_set], [k for k in traj_keys_set if k not in default_traj_keys]
+        assert traj_keys_set == set(
+            default_traj_keys
+        ), "Keys of traj dict did not match standard form.\nMissing keys: {}\nAdditional keys: {}".format(
+            [k for k in default_traj_keys if k not in traj_keys_set],
+            [k for k in traj_keys_set if k not in default_traj_keys],
         )
-    
+
     @staticmethod
     def _check_right_types(trajectories):
         for idx in range(len(trajectories["ep_states"])):
-            states, actions, rewards = trajectories["ep_states"][idx], trajectories["ep_actions"][idx], trajectories["ep_rewards"][idx]
-            mdp_params, env_params = trajectories["mdp_params"][idx], trajectories["env_params"][idx]
+            states, actions, rewards = (
+                trajectories["ep_states"][idx],
+                trajectories["ep_actions"][idx],
+                trajectories["ep_rewards"][idx],
+            )
+            mdp_params, env_params = (
+                trajectories["mdp_params"][idx],
+                trajectories["env_params"][idx],
+            )
             assert all(type(j_a) is tuple for j_a in actions)
             assert all(type(s) is OvercookedState for s in states)
             assert type(mdp_params) is dict
@@ -184,19 +364,32 @@ def _check_right_types(trajectories):
 
     @staticmethod
     def _check_trajectories_dynamics(trajectories, verbose=True):
-        if any(env_params["num_mdp"] > 1 for env_params in trajectories["env_params"]):
+        if any(
+            env_params["num_mdp"] > 1
+            for env_params in trajectories["env_params"]
+        ):
             if verbose:
-                print("Skipping trajectory consistency checking because MDP was recognized as variable. "
-                    "Trajectory consistency checking is not yet supported for variable MDPs.")
+                print(
+                    "Skipping trajectory consistency checking because MDP was recognized as variable. "
+                    "Trajectory consistency checking is not yet supported for variable MDPs."
+                )
             return
 
-        _, envs = AgentEvaluator.get_mdps_and_envs_from_trajectories(trajectories)
+        _, envs = AgentEvaluator.get_mdps_and_envs_from_trajectories(
+            trajectories
+        )
 
         for idx in range(len(trajectories["ep_states"])):
-            states, actions, rewards = trajectories["ep_states"][idx], trajectories["ep_actions"][idx], trajectories["ep_rewards"][idx]
+            states, actions, rewards = (
+                trajectories["ep_states"][idx],
+                trajectories["ep_actions"][idx],
+                trajectories["ep_rewards"][idx],
+            )
             simulation_env = envs[idx]
 
-            assert len(states) == len(actions) == len(rewards), "# states {}\t# actions {}\t# rewards {}".format(
+            assert (
+                len(states) == len(actions) == len(rewards)
+            ), "# states {}\t# actions {}\t# rewards {}".format(
                 len(states), len(actions), len(rewards)
             )
 
@@ -205,12 +398,20 @@ def _check_trajectories_dynamics(trajectories, verbose=True):
                 curr_state = states[i]
                 simulation_env.state = curr_state
 
-                next_state, reward, done, info = simulation_env.step(actions[i])
+                next_state, reward, done, info = simulation_env.step(
+                    actions[i]
+                )
 
-                assert states[i + 1] == next_state, "States differed (expected vs actual): {}\n\nexpected dict: \t{}\nactual dict: \t{}".format(
-                    simulation_env.display_states(states[i + 1], next_state), states[i+1].to_dict(), next_state.to_dict()
+                assert (
+                    states[i + 1] == next_state
+                ), "States differed (expected vs actual): {}\n\nexpected dict: \t{}\nactual dict: \t{}".format(
+                    simulation_env.display_states(states[i + 1], next_state),
+                    states[i + 1].to_dict(),
+                    next_state.to_dict(),
+                )
+                assert rewards[i] == reward, "{} \t {}".format(
+                    rewards[i], reward
                 )
-                assert rewards[i] == reward, "{} \t {}".format(rewards[i], reward)
 
     @staticmethod
     def get_mdps_and_envs_from_trajectories(trajectories):
@@ -224,15 +425,18 @@ def get_mdps_and_envs_from_trajectories(trajectories):
             envs.append(env)
         return mdps, envs
 
-
     ### I/O METHODS ###
 
     @staticmethod
     def save_trajectories(trajectories, filename):
         AgentEvaluator.check_trajectories(trajectories)
-        if any(t["env_params"]["start_state_fn"] is not None for t in trajectories):
-            print("Saving trajectories with a custom start state. This can currently "
-                  "cause things to break when loading in the trajectories.")
+        if any(
+            t["env_params"]["start_state_fn"] is not None for t in trajectories
+        ):
+            print(
+                "Saving trajectories with a custom start state. This can currently "
+                "cause things to break when loading in the trajectories."
+            )
         save_pickle(trajectories, filename)
 
     @staticmethod
@@ -244,9 +448,13 @@ def load_trajectories(filename):
     @staticmethod
     def save_traj_as_json(trajectory, filename):
         """Saves the `idx`th trajectory as a list of state action pairs"""
-        assert set(DEFAULT_TRAJ_KEYS) == set(trajectory.keys()), "{} vs\n{}".format(DEFAULT_TRAJ_KEYS, trajectory.keys())
+        assert set(DEFAULT_TRAJ_KEYS) == set(
+            trajectory.keys()
+        ), "{} vs\n{}".format(DEFAULT_TRAJ_KEYS, trajectory.keys())
         AgentEvaluator.check_trajectories(trajectory)
-        trajectory = AgentEvaluator.make_trajectories_json_serializable(trajectory)
+        trajectory = AgentEvaluator.make_trajectories_json_serializable(
+            trajectory
+        )
         save_as_json(trajectory, filename)
 
     @staticmethod
@@ -256,25 +464,41 @@ def make_trajectories_json_serializable(trajectories):
         This method converts all components of a trajectory to standard types.
         """
         dict_traj = copy.deepcopy(trajectories)
-        dict_traj["ep_states"] = [[ob.to_dict() for ob in one_ep_obs] for one_ep_obs in trajectories["ep_states"]]
+        dict_traj["ep_states"] = [
+            [ob.to_dict() for ob in one_ep_obs]
+            for one_ep_obs in trajectories["ep_states"]
+        ]
         for k in dict_traj.keys():
             dict_traj[k] = list(dict_traj[k])
-        dict_traj['ep_actions'] = [list(lst) for lst in dict_traj['ep_actions']]
-        dict_traj['ep_rewards'] = [list(lst) for lst in dict_traj['ep_rewards']]
-        dict_traj['ep_dones'] = [list(lst) for lst in dict_traj['ep_dones']]
-        dict_traj['ep_returns'] = [int(val) for val in dict_traj['ep_returns']]
-        dict_traj['ep_lengths'] = [int(val) for val in dict_traj['ep_lengths']]
+        dict_traj["ep_actions"] = [
+            list(lst) for lst in dict_traj["ep_actions"]
+        ]
+        dict_traj["ep_rewards"] = [
+            list(lst) for lst in dict_traj["ep_rewards"]
+        ]
+        dict_traj["ep_dones"] = [list(lst) for lst in dict_traj["ep_dones"]]
+        dict_traj["ep_returns"] = [int(val) for val in dict_traj["ep_returns"]]
+        dict_traj["ep_lengths"] = [int(val) for val in dict_traj["ep_lengths"]]
 
         # NOTE: Currently saving to JSON does not support ep_infos (due to nested np.arrays) or metadata
-        del dict_traj['ep_infos']
-        del dict_traj['metadatas']
+        del dict_traj["ep_infos"]
+        del dict_traj["metadatas"]
         return dict_traj
 
     @staticmethod
     def load_traj_from_json(filename):
         traj_dict = load_from_json(filename)
-        traj_dict["ep_states"] = [[OvercookedState.from_dict(ob) for ob in curr_ep_obs] for curr_ep_obs in traj_dict["ep_states"]]
-        traj_dict["ep_actions"] = [[tuple(tuple(a) if type(a) is list else a for a in j_a) for j_a in ep_acts] for ep_acts in traj_dict["ep_actions"]]
+        traj_dict["ep_states"] = [
+            [OvercookedState.from_dict(ob) for ob in curr_ep_obs]
+            for curr_ep_obs in traj_dict["ep_states"]
+        ]
+        traj_dict["ep_actions"] = [
+            [
+                tuple(tuple(a) if type(a) is list else a for a in j_a)
+                for j_a in ep_acts
+            ]
+            for ep_acts in traj_dict["ep_actions"]
+        ]
         return traj_dict
 
     ############################
@@ -289,7 +513,9 @@ def merge_trajs(trajs_n):
 
         [trajs0, trajs1] -> trajs
         """
-        metadatas_merged = merge_dictionaries([trajs["metadatas"] for trajs in trajs_n])
+        metadatas_merged = merge_dictionaries(
+            [trajs["metadatas"] for trajs in trajs_n]
+        )
         merged_trajs = merge_dictionaries(trajs_n)
         merged_trajs["metadatas"] = metadatas_merged
         return merged_trajs
@@ -308,15 +534,19 @@ def remove_traj_idx(trajs, idx):
     @staticmethod
     def take_traj_indices(trajs, indices):
         # NOTE: non mutating method
-        subset_trajs = take_indexes_from_dict(trajs, indices, keys_to_ignore=["metadatas"])
+        subset_trajs = take_indexes_from_dict(
+            trajs, indices, keys_to_ignore=["metadatas"]
+        )
         # TODO: Make metadatas field into additional keys for trajs, rather than having a metadatas field?
-        subset_trajs["metadatas"] = take_indexes_from_dict(trajs["metadatas"], indices)
+        subset_trajs["metadatas"] = take_indexes_from_dict(
+            trajs["metadatas"], indices
+        )
         return subset_trajs
 
     @staticmethod
     def add_metadata_to_traj(trajs, metadata_fn, input_keys):
         """
-        Add an additional metadata entry to the trajectory, based on manipulating 
+        Add an additional metadata entry to the trajectory, based on manipulating
         the trajectory `input_keys` values
         """
         metadata_fn_input = [trajs[k] for k in input_keys]
@@ -328,16 +558,20 @@ def add_metadata_to_traj(trajs, metadata_fn, input_keys):
     @staticmethod
     def add_observations_to_trajs_in_metadata(trajs, encoding_fn):
         """Adds processed observations (for both agent indices) in the metadatas"""
+
         def metadata_fn(data):
             traj_ep_states = data[0]
             obs_metadata = []
             for one_traj_states in traj_ep_states:
                 obs_metadata.append([encoding_fn(s) for s in one_traj_states])
             return "ep_obs_for_both_agents", obs_metadata
-        return AgentEvaluator.add_metadata_to_traj(trajs, metadata_fn, ["ep_states"])
+
+        return AgentEvaluator.add_metadata_to_traj(
+            trajs, metadata_fn, ["ep_states"]
+        )
 
     # EVENTS VISUALIZATION METHODS #
-    
+
     @staticmethod
     def events_visualization(trajs, traj_index):
         # TODO
diff --git a/src/overcooked_ai_py/data/planners/__init__.py b/src/overcooked_ai_py/data/planners/__init__.py
index d38f72cf..4e1dd800 100644
--- a/src/overcooked_ai_py/data/planners/__init__.py
+++ b/src/overcooked_ai_py/data/planners/__init__.py
@@ -1,14 +1,17 @@
-import os, pickle
-from overcooked_ai_py.utils import load_dict_from_file
+import os
+import pickle
+
 from overcooked_ai_py.static import PLANNERS_DIR
+from overcooked_ai_py.utils import load_dict_from_file
+
 
 def load_saved_action_manager(filename):
-    with open(os.path.join(PLANNERS_DIR, filename), 'rb') as f:
+    with open(os.path.join(PLANNERS_DIR, filename), "rb") as f:
         mlp_action_manager = pickle.load(f)
         return mlp_action_manager
 
 
 def load_saved_motion_planner(filename):
-    with open(os.path.join(PLANNERS_DIR, filename), 'rb') as f:
+    with open(os.path.join(PLANNERS_DIR, filename), "rb") as f:
         motion_planner = pickle.load(f)
-        return motion_planner
\ No newline at end of file
+        return motion_planner
diff --git a/src/overcooked_ai_py/mdp/actions.py b/src/overcooked_ai_py/mdp/actions.py
index 41fa1a34..02179c71 100644
--- a/src/overcooked_ai_py/mdp/actions.py
+++ b/src/overcooked_ai_py/mdp/actions.py
@@ -1,4 +1,6 @@
-import itertools, copy
+import copy
+import itertools
+
 import numpy as np
 
 
@@ -6,15 +8,20 @@ class Direction(object):
     """
     The four possible directions a player can be facing.
     """
-    
+
     NORTH = (0, -1)
     SOUTH = (0, 1)
-    EAST  = (1, 0)
-    WEST  = (-1, 0)
+    EAST = (1, 0)
+    WEST = (-1, 0)
     ALL_DIRECTIONS = INDEX_TO_DIRECTION = [NORTH, SOUTH, EAST, WEST]
-    DIRECTION_TO_INDEX = { a:i for i, a in enumerate(INDEX_TO_DIRECTION) }
-    OPPOSITE_DIRECTIONS = { NORTH: SOUTH, SOUTH: NORTH, EAST: WEST, WEST: EAST }
-    DIRECTION_TO_NAME = { d:name for d, name in zip([NORTH, SOUTH, EAST, WEST], ["NORTH", "SOUTH", "EAST", "WEST"])}
+    DIRECTION_TO_INDEX = {a: i for i, a in enumerate(INDEX_TO_DIRECTION)}
+    OPPOSITE_DIRECTIONS = {NORTH: SOUTH, SOUTH: NORTH, EAST: WEST, WEST: EAST}
+    DIRECTION_TO_NAME = {
+        d: name
+        for d, name in zip(
+            [NORTH, SOUTH, EAST, WEST], ["NORTH", "SOUTH", "EAST", "WEST"]
+        )
+    }
 
     @staticmethod
     def get_adjacent_directions(direction):
@@ -26,7 +33,7 @@ def get_adjacent_directions(direction):
             return [Direction.EAST, Direction.WEST]
         elif direction in [Direction.EAST, Direction.WEST]:
             return [Direction.NORTH, Direction.SOUTH]
-        raise ValueError('Invalid direction: %s' % direction)
+        raise ValueError("Invalid direction: %s" % direction)
 
 
 class Action(object):
@@ -38,18 +45,23 @@ class Action(object):
     """
 
     STAY = (0, 0)
-    INTERACT = 'interact'
-    ALL_ACTIONS = INDEX_TO_ACTION = Direction.INDEX_TO_DIRECTION + [STAY, INTERACT]
-    INDEX_TO_ACTION_INDEX_PAIRS = [v for v in itertools.product(range(len(INDEX_TO_ACTION)), repeat=2)]
-    ACTION_TO_INDEX = { a:i for i, a in enumerate(INDEX_TO_ACTION) }
+    INTERACT = "interact"
+    ALL_ACTIONS = INDEX_TO_ACTION = Direction.INDEX_TO_DIRECTION + [
+        STAY,
+        INTERACT,
+    ]
+    INDEX_TO_ACTION_INDEX_PAIRS = [
+        v for v in itertools.product(range(len(INDEX_TO_ACTION)), repeat=2)
+    ]
+    ACTION_TO_INDEX = {a: i for i, a in enumerate(INDEX_TO_ACTION)}
     MOTION_ACTIONS = Direction.ALL_DIRECTIONS + [STAY]
-    ACTION_TO_CHAR = { 
+    ACTION_TO_CHAR = {
         Direction.NORTH: "↑",
         Direction.SOUTH: "↓",
         Direction.EAST: "→",
         Direction.WEST: "←",
         STAY: "stay",
-        INTERACT: INTERACT
+        INTERACT: INTERACT,
     }
     NUM_ACTIONS = len(ALL_ACTIONS)
 
@@ -79,8 +91,10 @@ def determine_action_for_change_in_pos(old_pos, new_pos):
 
     @staticmethod
     def sample(action_probs):
-        return np.random.choice(np.array(Action.ALL_ACTIONS, dtype=object), p=action_probs)
-    
+        return np.random.choice(
+            np.array(Action.ALL_ACTIONS, dtype=object), p=action_probs
+        )
+
     @staticmethod
     def argmax(action_probs):
         action_idx = np.argmax(action_probs)
@@ -94,7 +108,7 @@ def remove_indices_and_renormalize(probs, indices, eps=0.0):
             for row_idx, row in enumerate(indices):
                 for idx in indices:
                     probs[row_idx][idx] = eps
-            norm_probs =  probs.T / np.sum(probs, axis=1)
+            norm_probs = probs.T / np.sum(probs, axis=1)
             return norm_probs.T
         else:
             for idx in indices:
@@ -114,4 +128,4 @@ def joint_action_to_char(joint_action):
     @staticmethod
     def uniform_probs_over_actions():
         num_acts = len(Action.ALL_ACTIONS)
-        return np.ones(num_acts) / num_acts
\ No newline at end of file
+        return np.ones(num_acts) / num_acts
diff --git a/src/overcooked_ai_py/mdp/layout_generator.py b/src/overcooked_ai_py/mdp/layout_generator.py
index 100b9e19..7789921e 100644
--- a/src/overcooked_ai_py/mdp/layout_generator.py
+++ b/src/overcooked_ai_py/mdp/layout_generator.py
@@ -1,24 +1,32 @@
+import copy
+import random
+
 import numpy as np
 
-import random, copy
-from overcooked_ai_py.utils import rnd_int_uniform, rnd_uniform
 from overcooked_ai_py.mdp.actions import Action, Direction
 from overcooked_ai_py.mdp.overcooked_mdp import OvercookedGridworld, Recipe
+from overcooked_ai_py.utils import rnd_int_uniform, rnd_uniform
 
-EMPTY = ' '
-COUNTER = 'X'
-ONION_DISPENSER = 'O'
-TOMATO_DISPENSER = 'T'
-POT = 'P'
-DISH_DISPENSER = 'D'
-SERVING_LOC = 'S'
-
-CODE_TO_TYPE = {0: EMPTY, 1: COUNTER, 2: ONION_DISPENSER, 3: TOMATO_DISPENSER, 4: POT, 5: DISH_DISPENSER,
-                6: SERVING_LOC}
+EMPTY = " "
+COUNTER = "X"
+ONION_DISPENSER = "O"
+TOMATO_DISPENSER = "T"
+POT = "P"
+DISH_DISPENSER = "D"
+SERVING_LOC = "S"
+
+CODE_TO_TYPE = {
+    0: EMPTY,
+    1: COUNTER,
+    2: ONION_DISPENSER,
+    3: TOMATO_DISPENSER,
+    4: POT,
+    5: DISH_DISPENSER,
+    6: SERVING_LOC,
+}
 TYPE_TO_CODE = {v: k for k, v in CODE_TO_TYPE.items()}
 
 
-
 def mdp_fn_random_choice(mdp_fn_choices):
     assert type(mdp_fn_choices) is list and len(mdp_fn_choices) > 0
     return random.choice(mdp_fn_choices)
@@ -34,12 +42,10 @@ def mdp_fn_random_choice(mdp_fn_choices):
     "inner_shape": (5, 4),
     "prop_empty": 0.95,
     "prop_feats": 0.1,
-    "start_all_orders" : [
-        { "ingredients" : ["onion", "onion", "onion"]}
-    ],
-    "recipe_values" : [20],
-    "recipe_times" : [20],
-    "display": False
+    "start_all_orders": [{"ingredients": ["onion", "onion", "onion"]}],
+    "recipe_values": [20],
+    "recipe_times": [20],
+    "display": False,
 }
 
 
@@ -48,23 +54,22 @@ def DEFAILT_PARAMS_SCHEDULE_FN(outside_information):
         "inner_shape": (5, 4),
         "prop_empty": 0.95,
         "prop_feats": 0.1,
-        "start_all_orders" : [
-        { "ingredients" : ["onion", "onion", "onion"]}
-        ],
-        "recipe_values" : [20],
-        "recipe_times" : [20],
-        "display": False
+        "start_all_orders": [{"ingredients": ["onion", "onion", "onion"]}],
+        "recipe_values": [20],
+        "recipe_times": [20],
+        "display": False,
     }
     return mdp_default_gen_params
 
 
 class MDPParamsGenerator(object):
-
     def __init__(self, params_schedule_fn):
         """
         params_schedule_fn (callable): the function to produce a set of mdp_params for a specific layout
         """
-        assert callable(params_schedule_fn), "params scheduling function must be a callable"
+        assert callable(
+            params_schedule_fn
+        ), "params scheduling function must be a callable"
         self.params_schedule_fn = params_schedule_fn
 
     @staticmethod
@@ -82,7 +87,14 @@ def generate(self, outside_information={}):
         mdp_params = self.params_schedule_fn(outside_information)
         return mdp_params
 
-DEFAULT_FEATURE_TYPES = (POT, ONION_DISPENSER, DISH_DISPENSER, SERVING_LOC) # NOTE: TOMATO_DISPENSER is disabled by default
+
+DEFAULT_FEATURE_TYPES = (
+    POT,
+    ONION_DISPENSER,
+    DISH_DISPENSER,
+    SERVING_LOC,
+)  # NOTE: TOMATO_DISPENSER is disabled by default
+
 
 class LayoutGenerator(object):
     # NOTE: This class hasn't been tested extensively.
@@ -97,7 +109,7 @@ def __init__(self, mdp_params_generator, outer_shape=(5, 4)):
 
     @staticmethod
     def mdp_gen_fn_from_dict(
-            mdp_params, outer_shape=None, mdp_params_schedule_fn=None
+        mdp_params, outer_shape=None, mdp_params_schedule_fn=None
     ):
         """
         mdp_params: one set of fixed mdp parameter used by the enviroment
@@ -113,12 +125,18 @@ def mdp_gen_fn_from_dict(
             # there is no schedule, we are using the same set of mdp_params all the time
             if mdp_params_schedule_fn is None:
                 assert mdp_params is not None
-                mdp_pg = MDPParamsGenerator.from_fixed_param(mdp_params_always=mdp_params)
+                mdp_pg = MDPParamsGenerator.from_fixed_param(
+                    mdp_params_always=mdp_params
+                )
             else:
-                assert mdp_params is None, "please remove the mdp_params from the variable, " \
-                                           "because mdp_params_schedule_fn exist and we will " \
-                                           "always use the schedule_fn if it exist"
-                mdp_pg = MDPParamsGenerator(params_schedule_fn=mdp_params_schedule_fn)
+                assert mdp_params is None, (
+                    "please remove the mdp_params from the variable, "
+                    "because mdp_params_schedule_fn exist and we will "
+                    "always use the schedule_fn if it exist"
+                )
+                mdp_pg = MDPParamsGenerator(
+                    params_schedule_fn=mdp_params_schedule_fn
+                )
             lg = LayoutGenerator(mdp_pg, outer_shape)
             mdp_fn = lg.generate_padded_mdp
         return mdp_fn
@@ -127,45 +145,75 @@ def generate_padded_mdp(self, outside_information={}):
         """
         Return a PADDED MDP with mdp params specified in self.mdp_params
         """
-        mdp_gen_params = self.mdp_params_generator.generate(outside_information)
+        mdp_gen_params = self.mdp_params_generator.generate(
+            outside_information
+        )
 
         outer_shape = self.outer_shape
-        if "layout_name" in mdp_gen_params.keys() and mdp_gen_params["layout_name"] is not None:
+        if (
+            "layout_name" in mdp_gen_params.keys()
+            and mdp_gen_params["layout_name"] is not None
+        ):
             mdp = OvercookedGridworld.from_layout_name(**mdp_gen_params)
             mdp_generator_fn = lambda: self.padded_mdp(mdp)
         else:
-            required_keys = ["inner_shape", "prop_empty", "prop_feats", "display"]
+            required_keys = [
+                "inner_shape",
+                "prop_empty",
+                "prop_feats",
+                "display",
+            ]
             # with generate_all_orders key start_all_orders will be generated inside make_new_layout method
-            if not mdp_gen_params.get("generate_all_orders"): 
+            if not mdp_gen_params.get("generate_all_orders"):
                 required_keys.append("start_all_orders")
-            missing_keys = [k for k in required_keys if k not in mdp_gen_params.keys()]
+            missing_keys = [
+                k for k in required_keys if k not in mdp_gen_params.keys()
+            ]
             if len(missing_keys) != 0:
                 print("missing keys dict", mdp_gen_params)
-            assert len(missing_keys) == 0, "These keys were missing from the mdp_params: {}".format(missing_keys)
+            assert (
+                len(missing_keys) == 0
+            ), "These keys were missing from the mdp_params: {}".format(
+                missing_keys
+            )
             inner_shape = mdp_gen_params["inner_shape"]
-            assert inner_shape[0] <= outer_shape[0] and inner_shape[1] <= outer_shape[1], \
-                "inner_shape cannot fit into the outershap"
-            layout_generator = LayoutGenerator(self.mdp_params_generator, outer_shape=self.outer_shape)
-            
+            assert (
+                inner_shape[0] <= outer_shape[0]
+                and inner_shape[1] <= outer_shape[1]
+            ), "inner_shape cannot fit into the outershap"
+            layout_generator = LayoutGenerator(
+                self.mdp_params_generator, outer_shape=self.outer_shape
+            )
+
             if "feature_types" not in mdp_gen_params:
                 mdp_gen_params["feature_types"] = DEFAULT_FEATURE_TYPES
-            
-            mdp_generator_fn = lambda: layout_generator.make_new_layout(mdp_gen_params)
+
+            mdp_generator_fn = lambda: layout_generator.make_new_layout(
+                mdp_gen_params
+            )
         return mdp_generator_fn()
-    
+
     @staticmethod
     def create_base_params(mdp_gen_params):
-        assert mdp_gen_params.get("start_all_orders") or mdp_gen_params.get("generate_all_orders")
-        mdp_gen_params = LayoutGenerator.add_generated_mdp_params_orders(mdp_gen_params)
-        recipe_params = {"start_all_orders": mdp_gen_params["start_all_orders"]}
+        assert mdp_gen_params.get("start_all_orders") or mdp_gen_params.get(
+            "generate_all_orders"
+        )
+        mdp_gen_params = LayoutGenerator.add_generated_mdp_params_orders(
+            mdp_gen_params
+        )
+        recipe_params = {
+            "start_all_orders": mdp_gen_params["start_all_orders"]
+        }
         if mdp_gen_params.get("start_bonus_orders"):
-            recipe_params["start_bonus_orders"] = mdp_gen_params["start_bonus_orders"]
+            recipe_params["start_bonus_orders"] = mdp_gen_params[
+                "start_bonus_orders"
+            ]
         if "recipe_values" in mdp_gen_params:
             recipe_params["recipe_values"] = mdp_gen_params["recipe_values"]
         if "recipe_times" in mdp_gen_params:
             recipe_params["recipe_times"] = mdp_gen_params["recipe_times"]
         return recipe_params
-        
+
     @staticmethod
     def add_generated_mdp_params_orders(mdp_params):
         """
@@ -174,11 +222,15 @@ def add_generated_mdp_params_orders(mdp_params):
         """
         mdp_params = copy.deepcopy(mdp_params)
         if mdp_params.get("generate_all_orders"):
-            all_orders_kwargs = copy.deepcopy(mdp_params["generate_all_orders"])
+            all_orders_kwargs = copy.deepcopy(
+                mdp_params["generate_all_orders"]
+            )
 
             if all_orders_kwargs.get("recipes"):
-                 all_orders_kwargs["recipes"] = [Recipe.from_dict(r) for r in all_orders_kwargs["recipes"]]
-        
+                all_orders_kwargs["recipes"] = [
+                    Recipe.from_dict(r) for r in all_orders_kwargs["recipes"]
+                ]
+
             all_recipes = Recipe.generate_random_recipes(**all_orders_kwargs)
             mdp_params["start_all_orders"] = [r.to_dict() for r in all_recipes]
         else:
@@ -186,13 +238,19 @@ def add_generated_mdp_params_orders(mdp_params):
             all_recipes = Recipe.ALL_RECIPES
 
         if mdp_params.get("generate_bonus_orders"):
-            bonus_orders_kwargs = copy.deepcopy(mdp_params["generate_bonus_orders"])
+            bonus_orders_kwargs = copy.deepcopy(
+                mdp_params["generate_bonus_orders"]
+            )
 
-            if not bonus_orders_kwargs.get("recipes"): 
+            if not bonus_orders_kwargs.get("recipes"):
                 bonus_orders_kwargs["recipes"] = all_recipes
 
-            bonus_recipes = Recipe.generate_random_recipes(**bonus_orders_kwargs)
-            mdp_params["start_bonus_orders"] = [r.to_dict() for r in bonus_recipes]
+            bonus_recipes = Recipe.generate_random_recipes(
+                **bonus_orders_kwargs
+            )
+            mdp_params["start_bonus_orders"] = [
+                r.to_dict() for r in bonus_recipes
+            ]
         return mdp_params
 
     def padded_mdp(self, mdp, display=False):
@@ -201,7 +259,9 @@ def padded_mdp(self, mdp, display=False):
         padded_grid = self.embed_grid(grid)
 
         start_positions = self.get_random_starting_positions(padded_grid)
-        mdp_grid = self.padded_grid_to_layout_grid(padded_grid, start_positions, display=display)
+        mdp_grid = self.padded_grid_to_layout_grid(
+            padded_grid, start_positions, display=display
+        )
         return OvercookedGridworld.from_grid(mdp_grid)
 
     def make_new_layout(self, mdp_gen_params):
@@ -211,20 +271,32 @@ def make_new_layout(self, mdp_gen_params):
             prop_features=mdp_gen_params["prop_feats"],
             base_param=LayoutGenerator.create_base_params(mdp_gen_params),
             feature_types=mdp_gen_params["feature_types"],
-            display=mdp_gen_params["display"]
-        )      
-
-    def make_disjoint_sets_layout(self, inner_shape, prop_empty, prop_features, base_param, feature_types=DEFAULT_FEATURE_TYPES, display=True):        
+            display=mdp_gen_params["display"],
+        )
+
+    def make_disjoint_sets_layout(
+        self,
+        inner_shape,
+        prop_empty,
+        prop_features,
+        base_param,
+        feature_types=DEFAULT_FEATURE_TYPES,
+        display=True,
+    ):
         grid = Grid(inner_shape)
         self.dig_space_with_disjoint_sets(grid, prop_empty)
         self.add_features(grid, prop_features, feature_types)
 
         padded_grid = self.embed_grid(grid)
         start_positions = self.get_random_starting_positions(padded_grid)
-        mdp_grid = self.padded_grid_to_layout_grid(padded_grid, start_positions, display=display)
+        mdp_grid = self.padded_grid_to_layout_grid(
+            padded_grid, start_positions, display=display
+        )
         return OvercookedGridworld.from_grid(mdp_grid, base_param)
 
-    def padded_grid_to_layout_grid(self, padded_grid, start_positions, display=False):
+    def padded_grid_to_layout_grid(
+        self, padded_grid, start_positions, display=False
+    ):
         if display:
             print("Generated layout")
             print(padded_grid)
@@ -258,7 +330,9 @@ def embed_grid(self, grid):
 
     def dig_space_with_disjoint_sets(self, grid, prop_empty):
         dsets = DisjointSets([])
-        while not (grid.proportion_empty() > prop_empty and dsets.num_sets == 1):
+        while not (
+            grid.proportion_empty() > prop_empty and dsets.num_sets == 1
+        ):
             valid_dig_location = False
             while not valid_dig_location:
                 loc = grid.get_random_interior_location()
@@ -290,9 +364,11 @@ def dig_space_with_fringe_expansion(self, grid, prop_empty=0.1):
                 if grid.is_valid_dig_location(location):
                     fringe.add(location)
 
-    def add_features(self, grid, prop_features=0, feature_types=DEFAULT_FEATURE_TYPES):
+    def add_features(
+        self, grid, prop_features=0, feature_types=DEFAULT_FEATURE_TYPES
+    ):
         """
-        Places one round of basic features and then adds random features 
+        Places one round of basic features and then adds random features
         until prop_features of valid locations are filled"""
 
         valid_locations = grid.valid_feature_locations()
@@ -321,7 +397,6 @@ def get_random_starting_positions(self, grid, divider_x=None):
 
 
 class Grid(object):
-
     def __init__(self, shape):
         assert len(shape) == 2, "Grid must be 2 dimensional"
         grid = (np.ones(shape) * TYPE_TO_CODE[COUNTER]).astype(np.int)
@@ -359,7 +434,9 @@ def change_location(self, location, feature_string):
     def proportion_empty(self):
         flattened_grid = self.mtx.flatten()
         num_eligible = len(flattened_grid) - 2 * sum(self.shape) + 4
-        num_empty = sum([1 for x in flattened_grid if x == TYPE_TO_CODE[EMPTY]])
+        num_empty = sum(
+            [1 for x in flattened_grid if x == TYPE_TO_CODE[EMPTY]]
+        )
         return float(num_empty) / num_eligible
 
     def get_near_locations(self, location):
@@ -383,7 +460,12 @@ def is_valid_dig_location(self, location):
             return False
 
         # If one of the edges of the map, or outside the map
-        if x <= 0 or y <= 0 or x >= self.shape[0] - 1 or y >= self.shape[1] - 1:
+        if (
+            x <= 0
+            or y <= 0
+            or x >= self.shape[0] - 1
+            or y >= self.shape[1] - 1
+        ):
             return False
         return True
 
@@ -408,7 +490,13 @@ def is_valid_feature_location(self, location):
             return False
 
         # If location is next to at least one empty square
-        if any([loc for loc in self.get_near_locations(location) if CODE_TO_TYPE[self.terrain_at_loc(loc)] == EMPTY]):
+        if any(
+            [
+                loc
+                for loc in self.get_near_locations(location)
+                if CODE_TO_TYPE[self.terrain_at_loc(loc)] == EMPTY
+            ]
+        ):
             return True
         else:
             return False
@@ -437,7 +525,9 @@ def convert_to_string(self):
                 column.append(CODE_TO_TYPE[self.mtx[x][y]])
             rows.append(column)
         string_grid = np.array(rows)
-        assert np.array_equal(string_grid.T.shape, self.shape), "{} vs {}".format(string_grid.shape, self.shape)
+        assert np.array_equal(
+            string_grid.T.shape, self.shape
+        ), "{} vs {}".format(string_grid.shape, self.shape)
         return string_grid
 
     def __repr__(self):
@@ -451,7 +541,6 @@ def __repr__(self):
 
 
 class Fringe(object):
-
     def __init__(self, grid):
         self.fringe_list = []
         self.distribution = []
@@ -464,13 +553,17 @@ def add(self, item):
 
     def pop(self):
         assert len(self.fringe_list) > 0
-        choice_idx = np.random.choice(len(self.fringe_list), p=self.distribution)
+        choice_idx = np.random.choice(
+            len(self.fringe_list), p=self.distribution
+        )
         removed_pos = self.fringe_list.pop(choice_idx)
         self.update_probs()
         return removed_pos
 
     def update_probs(self):
-        self.distribution = np.ones(len(self.fringe_list)) / len(self.fringe_list)
+        self.distribution = np.ones(len(self.fringe_list)) / len(
+            self.fringe_list
+        )
 
 
 class DisjointSets(object):
diff --git a/src/overcooked_ai_py/mdp/overcooked_env.py b/src/overcooked_ai_py/mdp/overcooked_env.py
index 778feb2b..b803db45 100644
--- a/src/overcooked_ai_py/mdp/overcooked_env.py
+++ b/src/overcooked_ai_py/mdp/overcooked_env.py
@@ -1,18 +1,31 @@
-import gym, tqdm
 import time
+
+import gym
 import numpy as np
-from overcooked_ai_py.utils import mean_and_std_err, append_dictionaries
-from overcooked_ai_py.mdp.actions import Action
-from overcooked_ai_py.mdp.overcooked_mdp import OvercookedGridworld, EVENT_TYPES
-from overcooked_ai_py.mdp.overcooked_trajectory import TIMESTEP_TRAJ_KEYS, EPISODE_TRAJ_KEYS, DEFAULT_TRAJ_KEYS
-from overcooked_ai_py.planning.planners import MediumLevelActionManager, MotionPlanner, NO_COUNTERS_PARAMS
+import tqdm
 
-DEFAULT_ENV_PARAMS = {
-    "horizon": 400
-}
+from overcooked_ai_py.mdp.actions import Action
+from overcooked_ai_py.mdp.overcooked_mdp import (
+    EVENT_TYPES,
+    OvercookedGridworld,
+)
+from overcooked_ai_py.mdp.overcooked_trajectory import (
+    DEFAULT_TRAJ_KEYS,
+    EPISODE_TRAJ_KEYS,
+    TIMESTEP_TRAJ_KEYS,
+)
+from overcooked_ai_py.planning.planners import (
+    NO_COUNTERS_PARAMS,
+    MediumLevelActionManager,
+    MotionPlanner,
+)
+from overcooked_ai_py.utils import append_dictionaries, mean_and_std_err
+
+DEFAULT_ENV_PARAMS = {"horizon": 400}
 
 MAX_HORIZON = 1e10
 
+
 class OvercookedEnv(object):
     """
     An environment wrapper for the OvercookedGridworld Markov Decision Process.
@@ -25,12 +38,20 @@ class OvercookedEnv(object):
     > env = OvercookedEnv.from_mdp(mdp, horizon=400)
     """
 
-
     #########################
     # INSTANTIATION METHODS #
     #########################
 
-    def __init__(self, mdp_generator_fn, start_state_fn=None, horizon=MAX_HORIZON, mlam_params=NO_COUNTERS_PARAMS, info_level=0, num_mdp=1, initial_info={}):
+    def __init__(
+        self,
+        mdp_generator_fn,
+        start_state_fn=None,
+        horizon=MAX_HORIZON,
+        mlam_params=NO_COUNTERS_PARAMS,
+        info_level=0,
+        num_mdp=1,
+        initial_info={},
+    ):
         """
         mdp_generator_fn (callable):    A no-argument function that returns a OvercookedGridworld instance
         start_state_fn (callable):      Function that returns start state for the MDP, called at each environment reset
@@ -43,9 +64,11 @@ def __init__(self, mdp_generator_fn, start_state_fn=None, horizon=MAX_HORIZON, m
         TODO: Potentially make changes based on this discussion
         https://github.com/HumanCompatibleAI/overcooked_ai/pull/22#discussion_r416786847
         """
-        assert callable(mdp_generator_fn),  "OvercookedEnv takes in a OvercookedGridworld generator function. " \
-                                            "If trying to instantiate directly from a OvercookedGridworld " \
-                                            "instance, use the OvercookedEnv.from_mdp method"
+        assert callable(mdp_generator_fn), (
+            "OvercookedEnv takes in a OvercookedGridworld generator function. "
+            "If trying to instantiate directly from a OvercookedGridworld "
+            "instance, use the OvercookedEnv.from_mdp method"
+        )
         self.num_mdp = num_mdp
         self.variable_mdp = num_mdp > 1
         self.mdp_generator_fn = mdp_generator_fn
@@ -57,16 +80,19 @@ def __init__(self, mdp_generator_fn, start_state_fn=None, horizon=MAX_HORIZON, m
         self.info_level = info_level
         self.reset(outside_info=initial_info)
         if self.horizon >= MAX_HORIZON and self.info_level > 0:
-            print("Environment has (near-)infinite horizon and no terminal states. \
-                Reduce info level of OvercookedEnv to not see this message.")
+            print(
+                "Environment has (near-)infinite horizon and no terminal states. \
+                Reduce info level of OvercookedEnv to not see this message."
+            )
 
     @property
     def mlam(self):
         if self._mlam is None:
             if self.info_level > 0:
                 print("Computing MediumLevelActionManager")
-            self._mlam = MediumLevelActionManager.from_pickle_or_compute(self.mdp, self.mlam_params,
-                                                                  force_compute=False)
+            self._mlam = MediumLevelActionManager.from_pickle_or_compute(
+                self.mdp, self.mlam_params, force_compute=False
+            )
         return self._mlam
 
     @property
@@ -77,12 +103,22 @@ def mp(self):
             else:
                 if self.info_level > 0:
                     print("Computing MotionPlanner")
-                self._mp = MotionPlanner.from_pickle_or_compute(self.mdp, self.mlam_params["counter_goals"],
-                                                                force_compute=False)
+                self._mp = MotionPlanner.from_pickle_or_compute(
+                    self.mdp,
+                    self.mlam_params["counter_goals"],
+                    force_compute=False,
+                )
         return self._mp
 
     @staticmethod
-    def from_mdp(mdp, start_state_fn=None, horizon=MAX_HORIZON, mlam_params=NO_COUNTERS_PARAMS, info_level=1, num_mdp=None):
+    def from_mdp(
+        mdp,
+        start_state_fn=None,
+        horizon=MAX_HORIZON,
+        mlam_params=NO_COUNTERS_PARAMS,
+        info_level=1,
+        num_mdp=None,
+    ):
         """
         Create an OvercookedEnv directly from a OvercookedGridworld mdp
         rather than a mdp generating function.
@@ -97,10 +133,9 @@ def from_mdp(mdp, start_state_fn=None, horizon=MAX_HORIZON, mlam_params=NO_COUNT
             horizon=horizon,
             mlam_params=mlam_params,
             info_level=info_level,
-            num_mdp=1
+            num_mdp=1,
         )
 
-
     #####################
     # BASIC CLASS UTILS #
     #####################
@@ -116,7 +151,7 @@ def env_params(self):
             "start_state_fn": self.start_state_fn,
             "horizon": self.horizon,
             "info_level": self.info_level,
-            "num_mdp": self.num_mdp
+            "num_mdp": self.num_mdp,
         }
 
     def copy(self):
@@ -126,10 +161,9 @@ def copy(self):
             start_state_fn=self.start_state_fn,
             horizon=self.horizon,
             info_level=self.info_level,
-            num_mdp=self.num_mdp
+            num_mdp=self.num_mdp,
         )
 
-
     #############################
     # ENV VISUALIZATION METHODS #
     #############################
@@ -148,37 +182,55 @@ def display_states(self, *states):
             print(self)
         self.state = old_state
 
-    def print_state_transition(self, a_t, r_t, env_info, fname=None, display_phi=False):
+    def print_state_transition(
+        self, a_t, r_t, env_info, fname=None, display_phi=False
+    ):
         """
         Terminal graphics visualization of a state transition.
         """
         # TODO: turn this into a "formatting action probs" function and add action symbols too
-        action_probs = [None if "action_probs" not in agent_info.keys() else list(agent_info["action_probs"]) for agent_info in env_info["agent_infos"]]
-
-        action_probs = [ None if player_action_probs is None else [round(p, 2) for p in player_action_probs[0]] for player_action_probs in action_probs ]
+        action_probs = [
+            None
+            if "action_probs" not in agent_info.keys()
+            else list(agent_info["action_probs"])
+            for agent_info in env_info["agent_infos"]
+        ]
+
+        action_probs = [
+            None
+            if player_action_probs is None
+            else [round(p, 2) for p in player_action_probs[0]]
+            for player_action_probs in action_probs
+        ]
 
         if display_phi:
-            state_potential_str = "\nState potential = " + str(env_info["phi_s_prime"]) + "\t"
-            potential_diff_str = "Δ potential = " + str(
-                0.99 * env_info["phi_s_prime"] - env_info["phi_s"]) + "\n"  # Assuming gamma 0.99
+            state_potential_str = (
+                "\nState potential = " + str(env_info["phi_s_prime"]) + "\t"
+            )
+            potential_diff_str = (
+                "Δ potential = "
+                + str(0.99 * env_info["phi_s_prime"] - env_info["phi_s"])
+                + "\n"
+            )  # Assuming gamma 0.99
         else:
             state_potential_str = ""
             potential_diff_str = ""
 
         output_string = "Timestep: {}\nJoint action taken: {} \t Reward: {} + shaping_factor * {}\nAction probs by index: {} {} {}\n{}\n".format(
-                    self.state.timestep,
-                    tuple(Action.ACTION_TO_CHAR[a] for a in a_t),
-                    r_t,
-                    env_info["shaped_r_by_agent"],
-                    action_probs,
-                    state_potential_str,
-                    potential_diff_str,
-                    self)
+            self.state.timestep,
+            tuple(Action.ACTION_TO_CHAR[a] for a in a_t),
+            r_t,
+            env_info["shaped_r_by_agent"],
+            action_probs,
+            state_potential_str,
+            potential_diff_str,
+            self,
+        )
 
         if fname is None:
             print(output_string)
         else:
-            f = open(fname, 'a')
+            f = open(fname, "a")
             print(output_string, file=f)
             f.close()
 
@@ -186,28 +238,34 @@ def print_state_transition(self, a_t, r_t, env_info, fname=None, display_phi=Fal
     # BASIC ENV LOGIC #
     ###################
 
-    def step(self, joint_action, joint_agent_action_info=None, display_phi=False):
+    def step(
+        self, joint_action, joint_agent_action_info=None, display_phi=False
+    ):
         """Performs a joint action, updating the environment state
         and providing a reward.
-        
+
         On being done, stats about the episode are added to info:
             ep_sparse_r: the environment sparse reward, given only at soup delivery
             ep_shaped_r: the component of the reward that is due to reward shaped (excluding sparse rewards)
             ep_length: length of rollout
         """
         assert not self.is_done()
-        if joint_agent_action_info is None: joint_agent_action_info = [{}, {}]
-        next_state, mdp_infos = self.mdp.get_state_transition(self.state, joint_action, display_phi, self.mp)
+        if joint_agent_action_info is None:
+            joint_agent_action_info = [{}, {}]
+        next_state, mdp_infos = self.mdp.get_state_transition(
+            self.state, joint_action, display_phi, self.mp
+        )
 
-        # Update game_stats 
+        # Update game_stats
         self._update_game_stats(mdp_infos)
 
         # Update state and done
         self.state = next_state
         done = self.is_done()
         env_info = self._prepare_info_dict(joint_agent_action_info, mdp_infos)
-        
-        if done: self._add_episode_info(env_info)
+
+        if done:
+            self._add_episode_info(env_info)
 
         timestep_sparse_reward = sum(mdp_infos["sparse_reward_by_agent"])
         return (next_state, timestep_sparse_reward, done, env_info)
@@ -244,16 +302,24 @@ def reset(self, regen_mdp=True, outside_info={}):
         else:
             self.state = self.start_state_fn()
 
-        events_dict = { k : [ [] for _ in range(self.mdp.num_players) ] for k in EVENT_TYPES }
+        events_dict = {
+            k: [[] for _ in range(self.mdp.num_players)] for k in EVENT_TYPES
+        }
         rewards_dict = {
-            "cumulative_sparse_rewards_by_agent": np.array([0] * self.mdp.num_players),
-            "cumulative_shaped_rewards_by_agent": np.array([0] * self.mdp.num_players)
+            "cumulative_sparse_rewards_by_agent": np.array(
+                [0] * self.mdp.num_players
+            ),
+            "cumulative_shaped_rewards_by_agent": np.array(
+                [0] * self.mdp.num_players
+            ),
         }
         self.game_stats = {**events_dict, **rewards_dict}
 
     def is_done(self):
         """Whether the episode is over."""
-        return self.state.timestep >= self.horizon or self.mdp.is_terminal(self.state)
+        return self.state.timestep >= self.horizon or self.mdp.is_terminal(
+            self.state
+        )
 
     def potential(self, mlam, state=None, gamma=0.99):
         """
@@ -274,22 +340,39 @@ def _prepare_info_dict(self, joint_agent_action_info, mdp_infos):
         """
         # Get the agent action info, that could contain info about action probs, or other
         # custom user defined information
-        env_info = {"agent_infos": [joint_agent_action_info[agent_idx] for agent_idx in range(self.mdp.num_players)]}
-        # TODO: This can be further simplified by having all the mdp_infos copied over to the env_infos automatically 
+        env_info = {
+            "agent_infos": [
+                joint_agent_action_info[agent_idx]
+                for agent_idx in range(self.mdp.num_players)
+            ]
+        }
+        # TODO: This can be further simplified by having all the mdp_infos copied over to the env_infos automatically
         env_info["sparse_r_by_agent"] = mdp_infos["sparse_reward_by_agent"]
         env_info["shaped_r_by_agent"] = mdp_infos["shaped_reward_by_agent"]
-        env_info["phi_s"] = mdp_infos["phi_s"] if "phi_s" in mdp_infos else None
-        env_info["phi_s_prime"] = mdp_infos["phi_s_prime"] if "phi_s_prime" in mdp_infos else None
+        env_info["phi_s"] = (
+            mdp_infos["phi_s"] if "phi_s" in mdp_infos else None
+        )
+        env_info["phi_s_prime"] = (
+            mdp_infos["phi_s_prime"] if "phi_s_prime" in mdp_infos else None
+        )
         return env_info
 
     def _add_episode_info(self, env_info):
         env_info["episode"] = {
             "ep_game_stats": self.game_stats,
-            "ep_sparse_r": sum(self.game_stats["cumulative_sparse_rewards_by_agent"]),
-            "ep_shaped_r": sum(self.game_stats["cumulative_shaped_rewards_by_agent"]),
-            "ep_sparse_r_by_agent": self.game_stats["cumulative_sparse_rewards_by_agent"],
-            "ep_shaped_r_by_agent": self.game_stats["cumulative_shaped_rewards_by_agent"],
-            "ep_length": self.state.timestep
+            "ep_sparse_r": sum(
+                self.game_stats["cumulative_sparse_rewards_by_agent"]
+            ),
+            "ep_shaped_r": sum(
+                self.game_stats["cumulative_shaped_rewards_by_agent"]
+            ),
+            "ep_sparse_r_by_agent": self.game_stats[
+                "cumulative_sparse_rewards_by_agent"
+            ],
+            "ep_shaped_r_by_agent": self.game_stats[
+                "cumulative_shaped_rewards_by_agent"
+            ],
+            "ep_length": self.state.timestep,
         }
         return env_info
 
@@ -298,49 +381,67 @@ def _update_game_stats(self, infos):
         Update the game stats dict based on the events of the current step
         NOTE: the timer ticks after events are logged, so there can be events from time 0 to time self.horizon - 1
         """
-        self.game_stats["cumulative_sparse_rewards_by_agent"] += np.array(infos["sparse_reward_by_agent"])
-        self.game_stats["cumulative_shaped_rewards_by_agent"] += np.array(infos["shaped_reward_by_agent"])
+        self.game_stats["cumulative_sparse_rewards_by_agent"] += np.array(
+            infos["sparse_reward_by_agent"]
+        )
+        self.game_stats["cumulative_shaped_rewards_by_agent"] += np.array(
+            infos["shaped_reward_by_agent"]
+        )
 
         for event_type, bool_list_by_agent in infos["event_infos"].items():
             # For each event type, store the timestep if it occurred
             event_occurred_by_idx = [int(x) for x in bool_list_by_agent]
             for idx, event_by_agent in enumerate(event_occurred_by_idx):
                 if event_by_agent:
-                    self.game_stats[event_type][idx].append(self.state.timestep)
-
+                    self.game_stats[event_type][idx].append(
+                        self.state.timestep
+                    )
 
     ####################
     # TRAJECTORY LOGIC #
     ####################
 
     def execute_plan(self, start_state, joint_action_plan, display=False):
-        """Executes action_plan (a list of joint actions) from a start 
+        """Executes action_plan (a list of joint actions) from a start
         state in the mdp and returns the resulting state."""
         self.state = start_state
         done = False
-        if display: print("Starting state\n{}".format(self))
+        if display:
+            print("Starting state\n{}".format(self))
         for joint_action in joint_action_plan:
             self.step(joint_action)
             done = self.is_done()
-            if display: print(self)
-            if done: break
+            if display:
+                print(self)
+            if done:
+                break
         successor_state = self.state
         self.reset(False)
         return successor_state, done
 
-    def run_agents(self, agent_pair, include_final_state=False, display=False, dir=None, display_phi=False, display_until=np.Inf):
+    def run_agents(
+        self,
+        agent_pair,
+        include_final_state=False,
+        display=False,
+        dir=None,
+        display_phi=False,
+        display_until=np.Inf,
+    ):
         """
         Trajectory returned will a list of state-action pairs (s_t, joint_a_t, r_t, done_t, info_t).
         """
-        assert self.state.timestep == 0, "Did not reset environment before running agents"
+        assert (
+            self.state.timestep == 0
+        ), "Did not reset environment before running agents"
         trajectory = []
         done = False
         # default is to not print to file
         fname = None
 
         if dir != None:
-            fname = dir + '/roll_out_' + str(time.time()) + '.txt'
-            f = open(fname, 'w+')
+            fname = dir + "/roll_out_" + str(time.time()) + ".txt"
+            f = open(fname, "w+")
             print(self, file=f)
             f.close()
         while not done:
@@ -358,23 +459,45 @@ def run_agents(self, agent_pair, include_final_state=False, display=False, dir=N
             if display and self.state.timestep < display_until:
                 self.print_state_transition(a_t, r_t, info, fname, display_phi)
 
-        assert len(trajectory) == self.state.timestep, "{} vs {}".format(len(trajectory), self.state.timestep)
+        assert len(trajectory) == self.state.timestep, "{} vs {}".format(
+            len(trajectory), self.state.timestep
+        )
 
         # Add final state
         if include_final_state:
             trajectory.append((s_tp1, (None, None), 0, True, None))
 
-        total_sparse = sum(self.game_stats["cumulative_sparse_rewards_by_agent"])
-        total_shaped = sum(self.game_stats["cumulative_shaped_rewards_by_agent"])
-        return np.array(trajectory, dtype=object), self.state.timestep, total_sparse, total_shaped
+        total_sparse = sum(
+            self.game_stats["cumulative_sparse_rewards_by_agent"]
+        )
+        total_shaped = sum(
+            self.game_stats["cumulative_shaped_rewards_by_agent"]
+        )
+        return (
+            np.array(trajectory, dtype=object),
+            self.state.timestep,
+            total_sparse,
+            total_shaped,
+        )
 
-    def get_rollouts(self, agent_pair, num_games, display=False, dir=None, final_state=False, display_phi=False,
-                     display_until=np.Inf, metadata_fn=None, metadata_info_fn=None, info=True):
-        """
-        Simulate `num_games` number rollouts with the current agent_pair and returns processed 
+    def get_rollouts(
+        self,
+        agent_pair,
+        num_games,
+        display=False,
+        dir=None,
+        final_state=False,
+        display_phi=False,
+        display_until=np.Inf,
+        metadata_fn=None,
+        metadata_info_fn=None,
+        info=True,
+    ):
+        """
+        Simulate `num_games` number rollouts with the current agent_pair and returns processed
         trajectories.
 
-        Returning excessive information to be able to convert trajectories to any required format 
+        Returning excessive information to be able to convert trajectories to any required format
         (baselines, stable_baselines, etc)
 
         metadata_fn returns some metadata information computed at the end of each trajectory based on
@@ -382,17 +505,40 @@ def get_rollouts(self, agent_pair, num_games, display=False, dir=None, final_sta
 
         NOTE: this is the standard trajectories format used throughout the codebase
         """
-        trajectories = { k:[] for k in DEFAULT_TRAJ_KEYS }
+        trajectories = {k: [] for k in DEFAULT_TRAJ_KEYS}
         metadata_fn = (lambda x: {}) if metadata_fn is None else metadata_fn
-        metadata_info_fn = (lambda x: "") if metadata_info_fn is None else metadata_info_fn
-        range_iterator = tqdm.trange(num_games, desc="", leave=True) if info else range(num_games)
+        metadata_info_fn = (
+            (lambda x: "") if metadata_info_fn is None else metadata_info_fn
+        )
+        range_iterator = (
+            tqdm.trange(num_games, desc="", leave=True)
+            if info
+            else range(num_games)
+        )
         for i in range_iterator:
             agent_pair.set_mdp(self.mdp)
 
-            rollout_info = self.run_agents(agent_pair, display=display, dir=dir, include_final_state=final_state,
-                                           display_phi=display_phi, display_until=display_until)
-            trajectory, time_taken, tot_rews_sparse, _tot_rews_shaped = rollout_info
-            obs, actions, rews, dones, infos = trajectory.T[0], trajectory.T[1], trajectory.T[2], trajectory.T[3], trajectory.T[4]
+            rollout_info = self.run_agents(
+                agent_pair,
+                display=display,
+                dir=dir,
+                include_final_state=final_state,
+                display_phi=display_phi,
+                display_until=display_until,
+            )
+            (
+                trajectory,
+                time_taken,
+                tot_rews_sparse,
+                _tot_rews_shaped,
+            ) = rollout_info
+            obs, actions, rews, dones, infos = (
+                trajectory.T[0],
+                trajectory.T[1],
+                trajectory.T[2],
+                trajectory.T[3],
+                trajectory.T[4],
+            )
             trajectories["ep_states"].append(obs)
             trajectories["ep_actions"].append(actions)
             trajectories["ep_rewards"].append(rews)
@@ -413,7 +559,11 @@ def get_rollouts(self, agent_pair, num_games, display=False, dir=None, final_sta
             if info:
                 mu, se = mean_and_std_err(trajectories["ep_returns"])
                 description = "Avg rew: {:.2f} (std: {:.2f}, se: {:.2f}); avg len: {:.2f}; ".format(
-                    mu, np.std(trajectories["ep_returns"]), se, np.mean(trajectories["ep_lengths"]))
+                    mu,
+                    np.std(trajectories["ep_returns"]),
+                    se,
+                    np.mean(trajectories["ep_lengths"]),
+                )
                 description += metadata_info_fn(trajectories["metadatas"])
                 range_iterator.set_description(description)
                 range_iterator.refresh()
@@ -422,14 +572,16 @@ def get_rollouts(self, agent_pair, num_games, display=False, dir=None, final_sta
         trajectories = {k: np.array(v) for k, v in trajectories.items()}
 
         # Merging all metadata dictionaries, assumes same keys throughout all
-        trajectories["metadatas"] = append_dictionaries(trajectories["metadatas"])
+        trajectories["metadatas"] = append_dictionaries(
+            trajectories["metadatas"]
+        )
 
         # TODO: should probably transfer check methods over to Env class
         from overcooked_ai_py.agents.benchmarking import AgentEvaluator
+
         AgentEvaluator.check_trajectories(trajectories, verbose=info)
         return trajectories
 
-
     ####################
     # TRAJECTORY UTILS #
     ####################
@@ -438,7 +590,9 @@ def get_rollouts(self, agent_pair, num_games, display=False, dir=None, final_sta
     def get_discounted_rewards(trajectories, gamma):
         rews = trajectories["ep_rewards"]
         horizon = rews.shape[1]
-        return OvercookedEnv._get_discounted_rewards_with_horizon(rews, gamma, horizon)
+        return OvercookedEnv._get_discounted_rewards_with_horizon(
+            rews, gamma, horizon
+        )
 
     @staticmethod
     def _get_discounted_rewards_with_horizon(rewards_matrix, gamma, horizon):
@@ -464,7 +618,9 @@ def get_agent_infos_for_trajectories(trajectories, agent_idx):
         agent_infos = []
         for traj_idx in range(len(trajectories["ep_lengths"])):
             ep_infos = trajectories["ep_infos"][traj_idx]
-            traj_agent_infos = [step_info["agent_infos"][agent_idx] for step_info in ep_infos]
+            traj_agent_infos = [
+                step_info["agent_infos"][agent_idx] for step_info in ep_infos
+            ]
 
             # Append all dictionaries together
             traj_agent_infos = append_dictionaries(traj_agent_infos)
@@ -487,10 +643,16 @@ def proportion_stuck_time(trajectories, agent_idx, stuck_time=3):
         for traj_idx in range(len(trajectories["ep_lengths"])):
             stuck_matrix.append([])
             obs = trajectories["ep_states"][traj_idx]
-            for traj_timestep in range(stuck_time, trajectories["ep_lengths"][traj_idx]):
+            for traj_timestep in range(
+                stuck_time, trajectories["ep_lengths"][traj_idx]
+            ):
                 if traj_timestep >= stuck_time:
-                    recent_states = obs[traj_timestep - stuck_time : traj_timestep + 1]
-                    recent_player_pos_and_or = [s.players[agent_idx].pos_and_or for s in recent_states]
+                    recent_states = obs[
+                        traj_timestep - stuck_time : traj_timestep + 1
+                    ]
+                    recent_player_pos_and_or = [
+                        s.players[agent_idx].pos_and_or for s in recent_states
+                    ]
 
                     if len({item for item in recent_player_pos_and_or}) == 1:
                         # If there is only one item in the last stuck_time steps, then we classify the agent as stuck
@@ -510,33 +672,36 @@ class Overcooked(gym.Env):
     NOTE: Observations returned are in a dictionary format with various information that is
      necessary to be able to handle the multi-agent nature of the environment. There are probably
      better ways to handle this, but we found this to work with minor modifications to OpenAI Baselines.
-    
+
     NOTE: The index of the main agent (as gym envs are 'single-agent') in the mdp is randomized at each reset
      of the environment, and is kept track of by the self.agent_idx attribute. This means that it is necessary
      to pass on this information in the output to know for which agent index featurizations should be made for
      other agents.
-    
+
     For example, say one is training A0 paired with A1, and A1 takes a custom state featurization.
     Then in the runner.py loop in OpenAI Baselines, we will get the lossless encodings of the state,
     and the true Overcooked state. When we encode the true state to feed to A1, we also need to know
     what agent index it has in the environment (as encodings will be index dependent).
     """
+
     env_name = "Overcooked-v0"
 
-    def custom_init(self, base_env, featurize_fn, baselines_reproducible=False):
+    def custom_init(
+        self, base_env, featurize_fn, baselines_reproducible=False
+    ):
         """
         base_env: OvercookedEnv
         featurize_fn(mdp, state): fn used to featurize states returned in the 'both_agent_obs' field
         """
         if baselines_reproducible:
             # NOTE:
-            # This will cause all agent indices to be chosen in sync across simulation 
+            # This will cause all agent indices to be chosen in sync across simulation
             # envs (for each update, all envs will have index 0 or index 1).
             # This is to prevent the randomness of choosing agent indexes
             # from leaking when using subprocess-vec-env in baselines (which
             # seeding does not reach) i.e. having different results for different
             # runs with the same seed.
-            # The effect of this should be negligible, as all other randomness is 
+            # The effect of this should be negligible, as all other randomness is
             # controlled by the actual run seeds
             np.random.seed(0)
 
@@ -556,15 +721,22 @@ def _setup_observation_space(self):
 
     def step(self, action):
         """
-        action: 
+        action:
             (agent with index self.agent_idx action, other agent action)
             is a tuple with the joint action of the primary and secondary agents in index format
-        
+
         returns:
             observation: formatted to be standard input for self.agent_idx's policy
         """
-        assert all(self.action_space.contains(a) for a in action), "%r (%s) invalid"%(action, type(action))
-        agent_action, other_agent_action = [Action.INDEX_TO_ACTION[a] for a in action]
+        assert all(
+            self.action_space.contains(a) for a in action
+        ), "%r (%s) invalid" % (
+            action,
+            type(action),
+        )
+        agent_action, other_agent_action = [
+            Action.INDEX_TO_ACTION[a] for a in action
+        ]
 
         if self.agent_idx == 0:
             joint_action = (agent_action, other_agent_action)
@@ -577,21 +749,23 @@ def step(self, action):
             both_agents_ob = (ob_p0, ob_p1)
         else:
             both_agents_ob = (ob_p1, ob_p0)
-        
+
         env_info["policy_agent_idx"] = self.agent_idx
 
         if "episode" in env_info.keys():
             env_info["episode"]["policy_agent_idx"] = self.agent_idx
 
-        obs = {"both_agent_obs": both_agents_ob,
-                "overcooked_state": next_state,
-                "other_agent_env_idx": 1 - self.agent_idx}
+        obs = {
+            "both_agent_obs": both_agents_ob,
+            "overcooked_state": next_state,
+            "other_agent_env_idx": 1 - self.agent_idx,
+        }
         return obs, reward, done, env_info
 
     def reset(self):
         """
         When training on individual maps, we want to randomize which agent is assigned to which
-        starting location, in order to make sure that the agents are trained to be able to 
+        starting location, in order to make sure that the agents are trained to be able to
         complete the task starting at either of the hardcoded positions.
 
         NOTE: a nicer way to do this would be to just randomize starting positions, and not
@@ -606,9 +780,11 @@ def reset(self):
             both_agents_ob = (ob_p0, ob_p1)
         else:
             both_agents_ob = (ob_p1, ob_p0)
-        return {"both_agent_obs": both_agents_ob, 
-                "overcooked_state": self.base_env.state, 
-                "other_agent_env_idx": 1 - self.agent_idx}
+        return {
+            "both_agent_obs": both_agents_ob,
+            "overcooked_state": self.base_env.state,
+            "other_agent_env_idx": 1 - self.agent_idx,
+        }
 
     def render(self, mode="human", close=False):
         pass
diff --git a/src/overcooked_ai_py/mdp/overcooked_mdp.py b/src/overcooked_ai_py/mdp/overcooked_mdp.py
index e961faeb..929c3f91 100644
--- a/src/overcooked_ai_py/mdp/overcooked_mdp.py
+++ b/src/overcooked_ai_py/mdp/overcooked_mdp.py
@@ -1,36 +1,61 @@
-import itertools, copy, warnings
-import numpy as np
+import copy
+import itertools
+import warnings
+from collections import Counter, defaultdict
 from functools import reduce
-from collections import defaultdict, Counter
-from overcooked_ai_py.utils import pos_distance, read_layout_dict, classproperty, OvercookedException
+
+import numpy as np
+
 from overcooked_ai_py.mdp.actions import Action, Direction
+from overcooked_ai_py.utils import (
+    OvercookedException,
+    classproperty,
+    pos_distance,
+    read_layout_dict,
+)
 
 
 class Recipe:
     MAX_NUM_INGREDIENTS = 3
 
-    TOMATO = 'tomato'
-    ONION = 'onion'
+    TOMATO = "tomato"
+    ONION = "onion"
     ALL_INGREDIENTS = [ONION, TOMATO]
 
     ALL_RECIPES_CACHE = {}
-    STR_REP = {'tomato': "†", 'onion': "ø"}
+    STR_REP = {"tomato": "†", "onion": "ø"}
 
     _computed = False
     _configured = False
     _conf = {}
-    
+
     def __new__(cls, ingredients):
         if not cls._configured:
-            raise OvercookedException("Recipe class must be configured before recipes can be created")
+            raise OvercookedException(
+                "Recipe class must be configured before recipes can be created"
+            )
         # Some basic argument verification
-        if not ingredients or not hasattr(ingredients, '__iter__') or len(ingredients) == 0:
-            raise ValueError("Invalid input recipe. Must be ingredients iterable with non-zero length")
+        if (
+            not ingredients
+            or not hasattr(ingredients, "__iter__")
+            or len(ingredients) == 0
+        ):
+            raise ValueError(
+                "Invalid input recipe. Must be ingredients iterable with non-zero length"
+            )
         for elem in ingredients:
             if not elem in cls.ALL_INGREDIENTS:
-                raise ValueError("Invalid ingredient: {0}. Recipe can only contain ingredients {1}".format(elem, cls.ALL_INGREDIENTS))
+                raise ValueError(
+                    "Invalid ingredient: {0}. Recipe can only contain ingredients {1}".format(
+                        elem, cls.ALL_INGREDIENTS
+                    )
+                )
         if not len(ingredients) <= cls.MAX_NUM_INGREDIENTS:
-            raise ValueError("Recipe of length {0} is invalid. Recipe can contain at most {1} ingredients".format(len(ingredients), cls.MAX_NUM_INGREDIENTS))
+            raise ValueError(
+                "Recipe of length {0} is invalid. Recipe can contain at most {1} ingredients".format(
+                    len(ingredients), cls.MAX_NUM_INGREDIENTS
+                )
+            )
         key = hash(tuple(sorted(ingredients)))
         if key in cls.ALL_RECIPES_CACHE:
             return cls.ALL_RECIPES_CACHE[key]
@@ -48,7 +73,9 @@ def __int__(self):
         num_onions = len([_ for _ in self.ingredients if _ == Recipe.ONION])
 
         mixed_mask = int(bool(num_tomatoes * num_onions))
-        mixed_shift = (Recipe.MAX_NUM_INGREDIENTS + 1)**len(Recipe.ALL_INGREDIENTS)
+        mixed_shift = (Recipe.MAX_NUM_INGREDIENTS + 1) ** len(
+            Recipe.ALL_INGREDIENTS
+        )
         encoding = num_onions + (Recipe.MAX_NUM_INGREDIENTS + 1) * num_tomatoes
 
         return mixed_mask * encoding * mixed_shift + encoding
@@ -91,7 +118,9 @@ def __deepcopy__(self, memo):
     @classmethod
     def _compute_all_recipes(cls):
         for i in range(cls.MAX_NUM_INGREDIENTS):
-            for ingredient_list in itertools.combinations_with_replacement(cls.ALL_INGREDIENTS, i + 1):
+            for ingredient_list in itertools.combinations_with_replacement(
+                cls.ALL_INGREDIENTS, i + 1
+            ):
                 cls(ingredient_list)
 
     @property
@@ -100,7 +129,9 @@ def ingredients(self):
 
     @ingredients.setter
     def ingredients(self, _):
-        raise AttributeError("Recpes are read-only. Do not modify instance attributes after creation")
+        raise AttributeError(
+            "Recpes are read-only. Do not modify instance attributes after creation"
+        )
 
     @property
     def value(self):
@@ -109,9 +140,24 @@ def value(self):
         if self._value_mapping and self in self._value_mapping:
             return self._value_mapping[self]
         if self._onion_value and self._tomato_value:
-            num_onions = len([ingredient for ingredient in self.ingredients if ingredient == self.ONION])
-            num_tomatoes = len([ingredient for ingredient in self.ingredients if ingredient == self.TOMATO])
-            return self._tomato_value * num_tomatoes + self._onion_value * num_onions
+            num_onions = len(
+                [
+                    ingredient
+                    for ingredient in self.ingredients
+                    if ingredient == self.ONION
+                ]
+            )
+            num_tomatoes = len(
+                [
+                    ingredient
+                    for ingredient in self.ingredients
+                    if ingredient == self.TOMATO
+                ]
+            )
+            return (
+                self._tomato_value * num_tomatoes
+                + self._onion_value * num_onions
+            )
         return 20
 
     @property
@@ -121,13 +167,28 @@ def time(self):
         if self._time_mapping and self in self._time_mapping:
             return self._time_mapping[self]
         if self._onion_time and self._tomato_time:
-            num_onions = len([ingredient for ingredient in self.ingredients if ingredient == self.ONION])
-            num_tomatoes = len([ingredient for ingredient in self.ingredients if ingredient == self.TOMATO])
-            return self._onion_time * num_onions + self._tomato_time * num_tomatoes
+            num_onions = len(
+                [
+                    ingredient
+                    for ingredient in self.ingredients
+                    if ingredient == self.ONION
+                ]
+            )
+            num_tomatoes = len(
+                [
+                    ingredient
+                    for ingredient in self.ingredients
+                    if ingredient == self.TOMATO
+                ]
+            )
+            return (
+                self._onion_time * num_onions
+                + self._tomato_time * num_tomatoes
+            )
         return 20
 
     def to_dict(self):
-        return { 'ingredients' : self.ingredients }
+        return {"ingredients": self.ingredients}
 
     def neighbors(self):
         """
@@ -161,7 +222,7 @@ def configure(cls, conf):
         cls._conf = conf
         cls._configured = True
         cls._computed = False
-        cls.MAX_NUM_INGREDIENTS = conf.get('max_num_ingredients', 3)
+        cls.MAX_NUM_INGREDIENTS = conf.get("max_num_ingredients", 3)
 
         cls._cook_time = None
         cls._delivery_reward = None
@@ -175,68 +236,115 @@ def configure(cls, conf):
         ## Basic checks for validity ##
 
         # Mutual Exclusion
-        if 'tomato_time' in conf and not 'onion_time' in conf or 'onion_time' in conf and not 'tomato_time' in conf:
-            raise ValueError("Must specify both 'onion_time' and 'tomato_time'")
-        if 'tomato_value' in conf and not 'onion_value' in conf or 'onion_value' in conf and not 'tomato_value' in conf:
-            raise ValueError("Must specify both 'onion_value' and 'tomato_value'")
-        if 'tomato_value' in conf and 'delivery_reward' in conf:
-            raise ValueError("'delivery_reward' incompatible with '<ingredient>_value'")
-        if 'tomato_value' in conf and 'recipe_values' in conf:
-            raise ValueError("'recipe_values' incompatible with '<ingredient>_value'")
-        if 'recipe_values' in conf and 'delivery_reward' in conf:
-            raise ValueError("'delivery_reward' incompatible with 'recipe_values'")
-        if 'tomato_time' in conf and 'cook_time' in conf:
-            raise ValueError("'cook_time' incompatible with '<ingredient>_time")
-        if 'tomato_time' in conf and 'recipe_times' in conf:
-            raise ValueError("'recipe_times' incompatible with '<ingredient>_time'")
-        if 'recipe_times' in conf and 'cook_time' in conf:
-            raise ValueError("'delivery_reward' incompatible with 'recipe_times'")
+        if (
+            "tomato_time" in conf
+            and not "onion_time" in conf
+            or "onion_time" in conf
+            and not "tomato_time" in conf
+        ):
+            raise ValueError(
+                "Must specify both 'onion_time' and 'tomato_time'"
+            )
+        if (
+            "tomato_value" in conf
+            and not "onion_value" in conf
+            or "onion_value" in conf
+            and not "tomato_value" in conf
+        ):
+            raise ValueError(
+                "Must specify both 'onion_value' and 'tomato_value'"
+            )
+        if "tomato_value" in conf and "delivery_reward" in conf:
+            raise ValueError(
+                "'delivery_reward' incompatible with '<ingredient>_value'"
+            )
+        if "tomato_value" in conf and "recipe_values" in conf:
+            raise ValueError(
+                "'recipe_values' incompatible with '<ingredient>_value'"
+            )
+        if "recipe_values" in conf and "delivery_reward" in conf:
+            raise ValueError(
+                "'delivery_reward' incompatible with 'recipe_values'"
+            )
+        if "tomato_time" in conf and "cook_time" in conf:
+            raise ValueError(
+                "'cook_time' incompatible with '<ingredient>_time"
+            )
+        if "tomato_time" in conf and "recipe_times" in conf:
+            raise ValueError(
+                "'recipe_times' incompatible with '<ingredient>_time'"
+            )
+        if "recipe_times" in conf and "cook_time" in conf:
+            raise ValueError(
+                "'delivery_reward' incompatible with 'recipe_times'"
+            )
 
         # recipe_ lists and orders compatibility
-        if 'recipe_values' in conf:
-            if not 'all_orders' in conf or not conf['all_orders']:
-                raise ValueError("Must specify 'all_orders' if 'recipe_values' specified")
-            if not len(conf['all_orders']) == len(conf['recipe_values']):
-                raise ValueError("Number of recipes in 'all_orders' must be the same as number in 'recipe_values")
-        if 'recipe_times' in conf:
-            if not 'all_orders' in conf or not conf['all_orders']:
-                raise ValueError("Must specify 'all_orders' if 'recipe_times' specified")
-            if not len(conf['all_orders']) == len(conf['recipe_times']):
-                raise ValueError("Number of recipes in 'all_orders' must be the same as number in 'recipe_times")
-
-        
+        if "recipe_values" in conf:
+            if not "all_orders" in conf or not conf["all_orders"]:
+                raise ValueError(
+                    "Must specify 'all_orders' if 'recipe_values' specified"
+                )
+            if not len(conf["all_orders"]) == len(conf["recipe_values"]):
+                raise ValueError(
+                    "Number of recipes in 'all_orders' must be the same as number in 'recipe_values"
+                )
+        if "recipe_times" in conf:
+            if not "all_orders" in conf or not conf["all_orders"]:
+                raise ValueError(
+                    "Must specify 'all_orders' if 'recipe_times' specified"
+                )
+            if not len(conf["all_orders"]) == len(conf["recipe_times"]):
+                raise ValueError(
+                    "Number of recipes in 'all_orders' must be the same as number in 'recipe_times"
+                )
+
         ## Conifgure ##
 
-        if 'cook_time' in conf:
-            cls._cook_time = conf['cook_time']
+        if "cook_time" in conf:
+            cls._cook_time = conf["cook_time"]
 
-        if 'delivery_reward' in conf:
-            cls._delivery_reward = conf['delivery_reward']
+        if "delivery_reward" in conf:
+            cls._delivery_reward = conf["delivery_reward"]
 
-        if 'recipe_values' in conf:
+        if "recipe_values" in conf:
             cls._value_mapping = {
-                cls.from_dict(recipe) : value for (recipe, value) in zip(conf['all_orders'], conf['recipe_values'])
+                cls.from_dict(recipe): value
+                for (recipe, value) in zip(
+                    conf["all_orders"], conf["recipe_values"]
+                )
             }
 
-        if 'recipe_times' in conf:
+        if "recipe_times" in conf:
             cls._time_mapping = {
-                cls.from_dict(recipe) : time for (recipe, time) in zip(conf['all_orders'], conf['recipe_times'])
+                cls.from_dict(recipe): time
+                for (recipe, time) in zip(
+                    conf["all_orders"], conf["recipe_times"]
+                )
             }
 
-        if 'tomato_time' in conf:
-            cls._tomato_time = conf['tomato_time']
+        if "tomato_time" in conf:
+            cls._tomato_time = conf["tomato_time"]
+
+        if "onion_time" in conf:
+            cls._onion_time = conf["onion_time"]
 
-        if 'onion_time' in conf:
-            cls._onion_time = conf['onion_time']
+        if "tomato_value" in conf:
+            cls._tomato_value = conf["tomato_value"]
 
-        if 'tomato_value' in conf:
-            cls._tomato_value = conf['tomato_value']
+        if "onion_value" in conf:
+            cls._onion_value = conf["onion_value"]
 
-        if 'onion_value' in conf:
-            cls._onion_value = conf['onion_value']
-    
     @classmethod
-    def generate_random_recipes(cls, n=1, min_size=2, max_size=3, ingredients=None, recipes=None, unique=True):
+    def generate_random_recipes(
+        cls,
+        n=1,
+        min_size=2,
+        max_size=3,
+        ingredients=None,
+        recipes=None,
+        unique=True,
+    ):
         """
         n (int): how many recipes generate
         min_size (int): min generated recipe size
@@ -245,28 +353,33 @@ def generate_random_recipes(cls, n=1, min_size=2, max_size=3, ingredients=None,
         recipes (list(Recipe)): list of recipes to choose from (default is cls.ALL_RECIPES)
         unique (bool): if all recipes are unique (without repeats)
         """
-        if recipes is None: recipes = cls.ALL_RECIPES
+        if recipes is None:
+            recipes = cls.ALL_RECIPES
 
         ingredients = set(ingredients or cls.ALL_INGREDIENTS)
-        choice_replace = not(unique)
+        choice_replace = not (unique)
 
         assert 1 <= min_size <= max_size <= cls.MAX_NUM_INGREDIENTS
-        assert all(ingredient in cls.ALL_INGREDIENTS for ingredient in ingredients)
+        assert all(
+            ingredient in cls.ALL_INGREDIENTS for ingredient in ingredients
+        )
 
         def valid_size(r):
             return min_size <= len(r.ingredients) <= max_size
 
         def valid_ingredients(r):
             return all(i in ingredients for i in r.ingredients)
-        
-        relevant_recipes = [r for r in recipes if valid_size(r) and valid_ingredients(r)]
+
+        relevant_recipes = [
+            r for r in recipes if valid_size(r) and valid_ingredients(r)
+        ]
         assert choice_replace or (n <= len(relevant_recipes))
         return np.random.choice(relevant_recipes, n, replace=choice_replace)
 
     @classmethod
     def from_dict(cls, obj_dict):
         return cls(**obj_dict)
-        
+
 
 class ObjectState(object):
     """
@@ -290,28 +403,26 @@ def position(self, new_pos):
         self._position = new_pos
 
     def is_valid(self):
-        return self.name in ['onion', 'tomato', 'dish']
+        return self.name in ["onion", "tomato", "dish"]
 
     def deepcopy(self):
         return ObjectState(self.name, self.position)
 
     def __eq__(self, other):
-        return isinstance(other, ObjectState) and \
-            self.name == other.name and \
-            self.position == other.position
+        return (
+            isinstance(other, ObjectState)
+            and self.name == other.name
+            and self.position == other.position
+        )
 
     def __hash__(self):
         return hash((self.name, self.position))
 
     def __repr__(self):
-        return '{}@{}'.format(
-            self.name, self.position)
+        return "{}@{}".format(self.name, self.position)
 
     def to_dict(self):
-        return {
-            "name": self.name,
-            "position": self.position
-        }
+        return {"name": self.name, "position": self.position}
 
     @classmethod
     def from_dict(cls, obj_dict):
@@ -320,12 +431,18 @@ def from_dict(cls, obj_dict):
 
 
 class SoupState(ObjectState):
-
-    def __init__(self, position, ingredients=[], cooking_tick=-1, cook_time=None, **kwargs):
+    def __init__(
+        self,
+        position,
+        ingredients=[],
+        cooking_tick=-1,
+        cook_time=None,
+        **kwargs
+    ):
         """
         Represents a soup object. An object becomes a soup the instant it is placed in a pot. The
         soup's recipe is a list of ingredient names used to create it. A soup's recipe is undetermined
-        until it has begun cooking. 
+        until it has begun cooking.
 
         position (tupe): (x, y) coordinates in the grid
         ingrdients (list(ObjectState)): Objects that have been used to cook this soup. Determiens @property recipe
@@ -339,8 +456,20 @@ def __init__(self, position, ingredients=[], cooking_tick=-1, cook_time=None, **
         self._cook_time = cook_time
 
     def __eq__(self, other):
-        return isinstance(other, SoupState) and self.name == other.name and self.position == other.position and self._cooking_tick == other._cooking_tick and \
-            all([this_i == other_i for this_i, other_i in zip(self._ingredients, other._ingredients)])
+        return (
+            isinstance(other, SoupState)
+            and self.name == other.name
+            and self.position == other.position
+            and self._cooking_tick == other._cooking_tick
+            and all(
+                [
+                    this_i == other_i
+                    for this_i, other_i in zip(
+                        self._ingredients, other._ingredients
+                    )
+                ]
+            )
+        )
 
     def __hash__(self):
         ingredient_hash = hash(tuple([hash(i) for i in self._ingredients]))
@@ -350,7 +479,9 @@ def __hash__(self):
     def __repr__(self):
         supercls_str = super(SoupState, self).__repr__()
         ingredients_str = self._ingredients.__repr__()
-        return "{}\nIngredients:\t{}\nCooking Tick:\t{}".format(supercls_str, ingredients_str, self._cooking_tick)
+        return "{}\nIngredients:\t{}\nCooking Tick:\t{}".format(
+            supercls_str, ingredients_str, self._cooking_tick
+        )
 
     def __str__(self):
         res = "{"
@@ -379,7 +510,9 @@ def is_cooking(self):
     @property
     def recipe(self):
         if self.is_idle:
-            raise ValueError("Recipe is not determined until soup begins cooking")
+            raise ValueError(
+                "Recipe is not determined until soup begins cooking"
+            )
         if not self._recipe:
             self._recipe = Recipe(self.ingredients)
         return self._recipe
@@ -412,10 +545,18 @@ def is_idle(self):
 
     @property
     def is_full(self):
-        return not self.is_idle or len(self.ingredients) == Recipe.MAX_NUM_INGREDIENTS
+        return (
+            not self.is_idle
+            or len(self.ingredients) == Recipe.MAX_NUM_INGREDIENTS
+        )
 
     def is_valid(self):
-        if not all([ingredient.position == self.position for ingredient in self._ingredients]):
+        if not all(
+            [
+                ingredient.position == self.position
+                for ingredient in self._ingredients
+            ]
+        ):
             return False
         if len(self.ingredients) > Recipe.MAX_NUM_INGREDIENTS:
             return False
@@ -441,7 +582,9 @@ def add_ingredient_from_str(self, ingredient_str):
 
     def pop_ingredient(self):
         if not self.is_idle:
-            raise ValueError("Cannot remove an ingredient from this soup at this time")
+            raise ValueError(
+                "Cannot remove an ingredient from this soup at this time"
+            )
         if len(self._ingredients) == 0:
             raise ValueError("No ingredient to remove")
         return self._ingredients.pop()
@@ -450,7 +593,9 @@ def begin_cooking(self):
         if not self.is_idle:
             raise ValueError("Cannot begin cooking this soup at this time")
         if len(self.ingredients) == 0:
-            raise ValueError("Must add at least one ingredient to soup before you can begin cooking")
+            raise ValueError(
+                "Must add at least one ingredient to soup before you can begin cooking"
+            )
         self._cooking_tick = 0
 
     def cook(self):
@@ -461,45 +606,72 @@ def cook(self):
         self._cooking_tick += 1
 
     def deepcopy(self):
-        return SoupState(self.position, [ingredient.deepcopy() for ingredient in self._ingredients], self._cooking_tick)
-    
+        return SoupState(
+            self.position,
+            [ingredient.deepcopy() for ingredient in self._ingredients],
+            self._cooking_tick,
+        )
+
     def to_dict(self):
         info_dict = super(SoupState, self).to_dict()
-        ingrdients_dict = [ingredient.to_dict() for ingredient in self._ingredients]
-        info_dict['_ingredients'] = ingrdients_dict
-        info_dict['cooking_tick'] = self._cooking_tick
-        info_dict['is_cooking'] = self.is_cooking
-        info_dict['is_ready'] = self.is_ready
-        info_dict['is_idle'] = self.is_idle
-        info_dict['cook_time'] = -1 if self.is_idle else self.cook_time
+        ingrdients_dict = [
+            ingredient.to_dict() for ingredient in self._ingredients
+        ]
+        info_dict["_ingredients"] = ingrdients_dict
+        info_dict["cooking_tick"] = self._cooking_tick
+        info_dict["is_cooking"] = self.is_cooking
+        info_dict["is_ready"] = self.is_ready
+        info_dict["is_idle"] = self.is_idle
+        info_dict["cook_time"] = -1 if self.is_idle else self.cook_time
 
         # This is for backwards compatibility w/ overcooked-demo
         # Should be removed once overcooked-demo is updated to use 'cooking_tick' instead of '_cooking_tick'
-        info_dict['_cooking_tick'] = self._cooking_tick
+        info_dict["_cooking_tick"] = self._cooking_tick
         return info_dict
 
     @classmethod
     def from_dict(cls, obj_dict):
         obj_dict = copy.deepcopy(obj_dict)
-        if obj_dict['name'] != 'soup':
+        if obj_dict["name"] != "soup":
             return super(SoupState, cls).from_dict(obj_dict)
 
-        if 'state' in obj_dict:
+        if "state" in obj_dict:
             # Legacy soup representation
-            ingredient, num_ingredient, time = obj_dict['state']
+            ingredient, num_ingredient, time = obj_dict["state"]
             cooking_tick = -1 if time == 0 else time
             finished = time >= 20
             if ingredient == Recipe.TOMATO:
-                return SoupState.get_soup(obj_dict['position'], num_tomatoes=num_ingredient, cooking_tick=cooking_tick, finished=finished)
+                return SoupState.get_soup(
+                    obj_dict["position"],
+                    num_tomatoes=num_ingredient,
+                    cooking_tick=cooking_tick,
+                    finished=finished,
+                )
             else:
-                return SoupState.get_soup(obj_dict['position'], num_onions=num_ingredient, cooking_tick=cooking_tick, finished=finished)
-
-        ingredients_objs = [ObjectState.from_dict(ing_dict) for ing_dict in obj_dict['_ingredients']]
-        obj_dict['ingredients'] = ingredients_objs
+                return SoupState.get_soup(
+                    obj_dict["position"],
+                    num_onions=num_ingredient,
+                    cooking_tick=cooking_tick,
+                    finished=finished,
+                )
+
+        ingredients_objs = [
+            ObjectState.from_dict(ing_dict)
+            for ing_dict in obj_dict["_ingredients"]
+        ]
+        obj_dict["ingredients"] = ingredients_objs
         return cls(**obj_dict)
 
     @classmethod
-    def get_soup(cls, position, num_onions=1, num_tomatoes=0, cooking_tick=-1, finished=False, **kwargs):
+    def get_soup(
+        cls,
+        position,
+        num_onions=1,
+        num_tomatoes=0,
+        cooking_tick=-1,
+        finished=False,
+        **kwargs
+    ):
         if num_onions < 0 or num_tomatoes < 0:
             raise ValueError("Number of active ingredients must be positive")
         if num_onions + num_tomatoes > Recipe.MAX_NUM_INGREDIENTS:
@@ -508,14 +680,18 @@ def get_soup(cls, position, num_onions=1, num_tomatoes=0, cooking_tick=-1, finis
             raise ValueError("_cooking_tick must be -1 for empty soup")
         if finished and num_tomatoes + num_onions == 0:
             raise ValueError("Empty soup cannot be finished")
-        onions = [ObjectState(Recipe.ONION, position) for _ in range(num_onions)]
-        tomatoes = [ObjectState(Recipe.TOMATO, position) for _ in range(num_tomatoes)]
+        onions = [
+            ObjectState(Recipe.ONION, position) for _ in range(num_onions)
+        ]
+        tomatoes = [
+            ObjectState(Recipe.TOMATO, position) for _ in range(num_tomatoes)
+        ]
         ingredients = onions + tomatoes
         soup = cls(position, ingredients, cooking_tick)
         if finished:
             soup.auto_finish()
         return soup
-        
+
 
 class PlayerState(object):
     """
@@ -526,6 +702,7 @@ class PlayerState(object):
     held_object: ObjectState representing the object held by the player, or
                  None if there is no such object.
     """
+
     def __init__(self, position, orientation, held_object=None):
         self.position = tuple(position)
         self.orientation = tuple(orientation)
@@ -551,13 +728,13 @@ def set_object(self, obj):
         assert not self.has_object()
         obj.position = self.position
         self.held_object = obj
- 
+
     def remove_object(self):
         assert self.has_object()
         obj = self.held_object
         self.held_object = None
         return obj
-    
+
     def update_pos_and_or(self, new_position, new_orientation):
         self.position = new_position
         self.orientation = new_orientation
@@ -565,27 +742,34 @@ def update_pos_and_or(self, new_position, new_orientation):
             self.get_object().position = new_position
 
     def deepcopy(self):
-        new_obj = None if self.held_object is None else self.held_object.deepcopy()
+        new_obj = (
+            None if self.held_object is None else self.held_object.deepcopy()
+        )
         return PlayerState(self.position, self.orientation, new_obj)
 
     def __eq__(self, other):
-        return isinstance(other, PlayerState) and \
-            self.position == other.position and \
-            self.orientation == other.orientation and \
-            self.held_object == other.held_object
+        return (
+            isinstance(other, PlayerState)
+            and self.position == other.position
+            and self.orientation == other.orientation
+            and self.held_object == other.held_object
+        )
 
     def __hash__(self):
         return hash((self.position, self.orientation, self.held_object))
 
     def __repr__(self):
-        return '{} facing {} holding {}'.format(
-            self.position, self.orientation, str(self.held_object))
-    
+        return "{} facing {} holding {}".format(
+            self.position, self.orientation, str(self.held_object)
+        )
+
     def to_dict(self):
         return {
             "position": self.position,
             "orientation": self.orientation,
-            "held_object": self.held_object.to_dict() if self.held_object is not None else None
+            "held_object": self.held_object.to_dict()
+            if self.held_object is not None
+            else None,
         }
 
     @staticmethod
@@ -599,11 +783,20 @@ def from_dict(player_dict):
 
 class OvercookedState(object):
     """A state in OvercookedGridworld."""
-    def __init__(self, players, objects, bonus_orders=[], all_orders=[], timestep=0, **kwargs):
+
+    def __init__(
+        self,
+        players,
+        objects,
+        bonus_orders=[],
+        all_orders=[],
+        timestep=0,
+        **kwargs
+    ):
         """
         players (list(PlayerState)): Currently active PlayerStates (index corresponds to number)
-        objects (dict({tuple:list(ObjectState)})):  Dictionary mapping positions (x, y) to ObjectStates. 
-            NOTE: Does NOT include objects held by players (they are in 
+        objects (dict({tuple:list(ObjectState)})):  Dictionary mapping positions (x, y) to ObjectStates.
+            NOTE: Does NOT include objects held by players (they are in
             the PlayerState objects).
         bonus_orders (list(dict)):   Current orders worth a bonus
         all_orders (list(dict)):     Current orders allowed at all
@@ -620,9 +813,15 @@ def __init__(self, players, objects, bonus_orders=[], all_orders=[], timestep=0,
         self._all_orders = all_orders
         self.timestep = timestep
 
-        assert len(set(self.bonus_orders)) == len(self.bonus_orders), "Bonus orders must not have duplicates"
-        assert len(set(self.all_orders)) == len(self.all_orders), "All orders must not have duplicates"
-        assert set(self.bonus_orders).issubset(set(self.all_orders)), "Bonus orders must be a subset of all orders"
+        assert len(set(self.bonus_orders)) == len(
+            self.bonus_orders
+        ), "Bonus orders must not have duplicates"
+        assert len(set(self.all_orders)) == len(
+            self.all_orders
+        ), "All orders must not have duplicates"
+        assert set(self.bonus_orders).issubset(
+            set(self.all_orders)
+        ), "Bonus orders must be a subset of all orders"
 
     @property
     def player_positions(self):
@@ -681,7 +880,11 @@ def all_objects_list(self):
 
     @property
     def all_orders(self):
-        return sorted(self._all_orders) if self._all_orders else sorted(Recipe.ALL_RECIPES)
+        return (
+            sorted(self._all_orders)
+            if self._all_orders
+            else sorted(Recipe.ALL_RECIPES)
+        )
 
     @property
     def bonus_orders(self):
@@ -715,71 +918,100 @@ def reverse_players(self):
         self.players = tuple(reversed)
         return self
 
-
     @classmethod
-    def from_players_pos_and_or(cls, players_pos_and_or, bonus_orders=[], all_orders=[]):
+    def from_players_pos_and_or(
+        cls, players_pos_and_or, bonus_orders=[], all_orders=[]
+    ):
         """
         Make a dummy OvercookedState with no objects based on the passed in player
         positions and orientations and order list
         """
         return cls(
-            [PlayerState(*player_pos_and_or) for player_pos_and_or in players_pos_and_or], 
-            objects={}, bonus_orders=bonus_orders, all_orders=all_orders)
+            [
+                PlayerState(*player_pos_and_or)
+                for player_pos_and_or in players_pos_and_or
+            ],
+            objects={},
+            bonus_orders=bonus_orders,
+            all_orders=all_orders,
+        )
 
     @classmethod
-    def from_player_positions(cls, player_positions, bonus_orders=[], all_orders=[]):
+    def from_player_positions(
+        cls, player_positions, bonus_orders=[], all_orders=[]
+    ):
         """
         Make a dummy OvercookedState with no objects and with players facing
         North based on the passed in player positions and order list
         """
         dummy_pos_and_or = [(pos, Direction.NORTH) for pos in player_positions]
-        return cls.from_players_pos_and_or(dummy_pos_and_or, bonus_orders, all_orders)
+        return cls.from_players_pos_and_or(
+            dummy_pos_and_or, bonus_orders, all_orders
+        )
 
     def deepcopy(self):
         return OvercookedState(
             players=[player.deepcopy() for player in self.players],
-            objects={pos:obj.deepcopy() for pos, obj in self.objects.items()}, 
+            objects={pos: obj.deepcopy() for pos, obj in self.objects.items()},
             bonus_orders=[order.to_dict() for order in self.bonus_orders],
             all_orders=[order.to_dict() for order in self.all_orders],
-            timestep=self.timestep)
+            timestep=self.timestep,
+        )
 
     def time_independent_equal(self, other):
-        order_lists_equal = self.all_orders == other.all_orders and self.bonus_orders == other.bonus_orders
+        order_lists_equal = (
+            self.all_orders == other.all_orders
+            and self.bonus_orders == other.bonus_orders
+        )
 
-        return isinstance(other, OvercookedState) and \
-            self.players == other.players and \
-            set(self.objects.items()) == set(other.objects.items()) and \
-            order_lists_equal
+        return (
+            isinstance(other, OvercookedState)
+            and self.players == other.players
+            and set(self.objects.items()) == set(other.objects.items())
+            and order_lists_equal
+        )
 
     def __eq__(self, other):
-        return self.time_independent_equal(other) and self.timestep == other.timestep
+        return (
+            self.time_independent_equal(other)
+            and self.timestep == other.timestep
+        )
 
     def __hash__(self):
         # NOTE: hash doesn't take into account timestep
-        order_list_hash = hash(tuple(self.bonus_orders)) + hash(tuple(self.all_orders))
+        order_list_hash = hash(tuple(self.bonus_orders)) + hash(
+            tuple(self.all_orders)
+        )
         return hash(
             (self.players, tuple(self.objects.values()), order_list_hash)
         )
 
     def __str__(self):
-        return 'Players: {}, Objects: {}, Bonus orders: {} All orders: {} Timestep: {}'.format( 
-            str(self.players), str(list(self.objects.values())), str(self.bonus_orders), str(self.all_orders), str(self.timestep))
+        return "Players: {}, Objects: {}, Bonus orders: {} All orders: {} Timestep: {}".format(
+            str(self.players),
+            str(list(self.objects.values())),
+            str(self.bonus_orders),
+            str(self.all_orders),
+            str(self.timestep),
+        )
 
     def to_dict(self):
         return {
             "players": [p.to_dict() for p in self.players],
             "objects": [obj.to_dict() for obj in self.objects.values()],
             "bonus_orders": [order.to_dict() for order in self.bonus_orders],
-            "all_orders" : [order.to_dict() for order in self.all_orders],
-            "timestep" : self.timestep
+            "all_orders": [order.to_dict() for order in self.all_orders],
+            "timestep": self.timestep,
         }
 
     @staticmethod
     def from_dict(state_dict):
         state_dict = copy.deepcopy(state_dict)
-        state_dict["players"] = [PlayerState.from_dict(p) for p in state_dict["players"]]
+        state_dict["players"] = [
+            PlayerState.from_dict(p) for p in state_dict["players"]
+        ]
         object_list = [SoupState.from_dict(o) for o in state_dict["objects"]]
-        state_dict["objects"] = { ob.position : ob for ob in object_list }
+        state_dict["objects"] = {ob.position: ob for ob in object_list}
         return OvercookedState(**state_dict)
 
 
@@ -789,61 +1021,58 @@ def from_dict(state_dict):
     "SOUP_PICKUP_REWARD": 5,
     "DISH_DISP_DISTANCE_REW": 0,
     "POT_DISTANCE_REW": 0,
-    "SOUP_DISTANCE_REW": 0
+    "SOUP_DISTANCE_REW": 0,
 }
 
 EVENT_TYPES = [
     # Tomato events
-    'tomato_pickup',
-    'useful_tomato_pickup',
-    'tomato_drop',
-    'useful_tomato_drop',
-    'potting_tomato',
-
+    "tomato_pickup",
+    "useful_tomato_pickup",
+    "tomato_drop",
+    "useful_tomato_drop",
+    "potting_tomato",
     # Onion events
-    'onion_pickup',
-    'useful_onion_pickup',
-    'onion_drop',
-    'useful_onion_drop',
-    'potting_onion',
-
+    "onion_pickup",
+    "useful_onion_pickup",
+    "onion_drop",
+    "useful_onion_drop",
+    "potting_onion",
     # Dish events
-    'dish_pickup',
-    'useful_dish_pickup',
-    'dish_drop',
-    'useful_dish_drop',
-
+    "dish_pickup",
+    "useful_dish_pickup",
+    "dish_drop",
+    "useful_dish_drop",
     # Soup events
-    'soup_pickup',
-    'soup_delivery',
-    'soup_drop',
-
+    "soup_pickup",
+    "soup_delivery",
+    "soup_drop",
     # Potting events
-    'optimal_onion_potting',
-    'optimal_tomato_potting',
-    'viable_onion_potting',
-    'viable_tomato_potting',
-    'catastrophic_onion_potting',
-    'catastrophic_tomato_potting',
-    'useless_onion_potting',
-    'useless_tomato_potting'
+    "optimal_onion_potting",
+    "optimal_tomato_potting",
+    "viable_onion_potting",
+    "viable_tomato_potting",
+    "catastrophic_onion_potting",
+    "catastrophic_tomato_potting",
+    "useless_onion_potting",
+    "useless_tomato_potting",
 ]
 
 POTENTIAL_CONSTANTS = {
-    'default' : {
-        'max_delivery_steps' : 10,
-        'max_pickup_steps' : 10,
-        'pot_onion_steps' : 10,
-        'pot_tomato_steps' : 10
+    "default": {
+        "max_delivery_steps": 10,
+        "max_pickup_steps": 10,
+        "pot_onion_steps": 10,
+        "pot_tomato_steps": 10,
+    },
+    "mdp_test_tomato": {
+        "max_delivery_steps": 4,
+        "max_pickup_steps": 4,
+        "pot_onion_steps": 5,
+        "pot_tomato_steps": 6,
     },
-    'mdp_test_tomato' : {
-        'max_delivery_steps' : 4,
-        'max_pickup_steps' : 4,
-        'pot_onion_steps' : 5,
-        'pot_tomato_steps' : 6
-    }
 }
 
+
 class OvercookedGridworld(object):
     """
     An MDP grid world based off of the Overcooked game.
@@ -854,17 +1083,28 @@ class OvercookedGridworld(object):
     TODO: clean the organization of this class further.
     """
 
-
     #########################
     # INSTANTIATION METHODS #
     #########################
 
-    def __init__(self, terrain, start_player_positions, start_bonus_orders=[], rew_shaping_params=None, layout_name="unnamed_layout", start_all_orders=[], num_items_for_soup=3, order_bonus=2, start_state=None, **kwargs):
+    def __init__(
+        self,
+        terrain,
+        start_player_positions,
+        start_bonus_orders=[],
+        rew_shaping_params=None,
+        layout_name="unnamed_layout",
+        start_all_orders=[],
+        num_items_for_soup=3,
+        order_bonus=2,
+        start_state=None,
+        **kwargs
+    ):
         """
         terrain: a matrix of strings that encode the MDP layout
         layout_name: string identifier of the layout
         start_player_positions: tuple of positions for both players' starting positions
-        start_bonus_orders: List of recipes dicts that are worth a bonus 
+        start_bonus_orders: List of recipes dicts that are worth a bonus
         rew_shaping_params: reward given for completion of specific subgoals
         all_orders: List of all available order dicts the players can make, defaults to all possible recipes if empy list provided
         num_items_for_soup: Maximum number of ingredients that can be placed in a soup
@@ -872,7 +1112,11 @@ def __init__(self, terrain, start_player_positions, start_bonus_orders=[], rew_s
         start_state: Default start state returned by get_standard_start_state
         """
         self._configure_recipes(start_all_orders, num_items_for_soup, **kwargs)
-        self.start_all_orders = [r.to_dict() for r in Recipe.ALL_RECIPES] if not start_all_orders else start_all_orders
+        self.start_all_orders = (
+            [r.to_dict() for r in Recipe.ALL_RECIPES]
+            if not start_all_orders
+            else start_all_orders
+        )
         self.height = len(terrain)
         self.width = len(terrain[0])
         self.shape = (self.width, self.height)
@@ -881,7 +1125,11 @@ def __init__(self, terrain, start_player_positions, start_bonus_orders=[], rew_s
         self.start_player_positions = start_player_positions
         self.num_players = len(start_player_positions)
         self.start_bonus_orders = start_bonus_orders
-        self.reward_shaping_params = BASE_REW_SHAPING_PARAMS if rew_shaping_params is None else rew_shaping_params
+        self.reward_shaping_params = (
+            BASE_REW_SHAPING_PARAMS
+            if rew_shaping_params is None
+            else rew_shaping_params
+        )
         self.layout_name = layout_name
         self.order_bonus = order_bonus
         self.start_state = start_state
@@ -889,7 +1137,6 @@ def __init__(self, terrain, start_player_positions, start_bonus_orders=[], rew_s
         self._opt_recipe_cache = {}
         self._prev_potential_params = {}
 
-
     @staticmethod
     def from_layout_name(layout_name, **params_to_overwrite):
         """
@@ -900,20 +1147,26 @@ def from_layout_name(layout_name, **params_to_overwrite):
         params_to_overwrite = params_to_overwrite.copy()
         base_layout_params = read_layout_dict(layout_name)
 
-        grid = base_layout_params['grid']
-        del base_layout_params['grid']
-        base_layout_params['layout_name'] = layout_name
-        if 'start_state' in base_layout_params:
-            base_layout_params['start_state'] = OvercookedState.from_dict(base_layout_params['start_state'])
+        grid = base_layout_params["grid"]
+        del base_layout_params["grid"]
+        base_layout_params["layout_name"] = layout_name
+        if "start_state" in base_layout_params:
+            base_layout_params["start_state"] = OvercookedState.from_dict(
+                base_layout_params["start_state"]
+            )
 
         # Clean grid
         grid = [layout_row.strip() for layout_row in grid.split("\n")]
-        return OvercookedGridworld.from_grid(grid, base_layout_params, params_to_overwrite)
+        return OvercookedGridworld.from_grid(
+            grid, base_layout_params, params_to_overwrite
+        )
 
     @staticmethod
-    def from_grid(layout_grid, base_layout_params={}, params_to_overwrite={}, debug=False):
+    def from_grid(
+        layout_grid, base_layout_params={}, params_to_overwrite={}, debug=False
+    ):
         """
-        Returns instance of OvercookedGridworld with terrain and starting 
+        Returns instance of OvercookedGridworld with terrain and starting
         positions derived from layout_grid.
         One can override default configuration parameters of the mdp in
         partial_mdp_config.
@@ -930,11 +1183,13 @@ def from_grid(layout_grid, base_layout_params={}, params_to_overwrite={}, debug=
         player_positions = [None] * 9
         for y, row in enumerate(layout_grid):
             for x, c in enumerate(row):
-                if c in ['1', '2', '3', '4', '5', '6', '7', '8', '9']:
-                    layout_grid[y][x] = ' '
+                if c in ["1", "2", "3", "4", "5", "6", "7", "8", "9"]:
+                    layout_grid[y][x] = " "
 
                     # -1 is to account for fact that player indexing starts from 1 rather than 0
-                    assert player_positions[int(c) - 1] is None, 'Duplicate player in grid'
+                    assert (
+                        player_positions[int(c) - 1] is None
+                    ), "Duplicate player in grid"
                     player_positions[int(c) - 1] = (x, y)
 
         num_players = len([x for x in player_positions if x is not None])
@@ -947,16 +1202,22 @@ def from_grid(layout_grid, base_layout_params={}, params_to_overwrite={}, debug=
         for k, v in params_to_overwrite.items():
             curr_val = mdp_config.get(k, None)
             if debug:
-                print("Overwriting mdp layout standard config value {}:{} -> {}".format(k, curr_val, v))
+                print(
+                    "Overwriting mdp layout standard config value {}:{} -> {}".format(
+                        k, curr_val, v
+                    )
+                )
             mdp_config[k] = v
 
         return OvercookedGridworld(**mdp_config)
 
-    def _configure_recipes(self, start_all_orders, num_items_for_soup, **kwargs):
+    def _configure_recipes(
+        self, start_all_orders, num_items_for_soup, **kwargs
+    ):
         self.recipe_config = {
-            "num_items_for_soup" : num_items_for_soup,
-            "all_orders" : start_all_orders,
-            **kwargs
+            "num_items_for_soup": num_items_for_soup,
+            "all_orders": start_all_orders,
+            **kwargs,
         }
         Recipe.configure(self.recipe_config)
 
@@ -965,13 +1226,15 @@ def _configure_recipes(self, start_all_orders, num_items_for_soup, **kwargs):
     #####################
 
     def __eq__(self, other):
-        return np.array_equal(self.terrain_mtx, other.terrain_mtx) and \
-                self.start_player_positions == other.start_player_positions and \
-                self.start_bonus_orders == other.start_bonus_orders and \
-                self.start_all_orders == other.start_all_orders and \
-                self.reward_shaping_params == other.reward_shaping_params and \
-                self.layout_name == other.layout_name
-    
+        return (
+            np.array_equal(self.terrain_mtx, other.terrain_mtx)
+            and self.start_player_positions == other.start_player_positions
+            and self.start_bonus_orders == other.start_bonus_orders
+            and self.start_all_orders == other.start_all_orders
+            and self.reward_shaping_params == other.reward_shaping_params
+            and self.layout_name == other.layout_name
+        )
+
     def copy(self):
         return OvercookedGridworld(
             terrain=self.terrain_mtx.copy(),
@@ -979,7 +1242,7 @@ def copy(self):
             start_bonus_orders=self.start_bonus_orders,
             rew_shaping_params=copy.deepcopy(self.reward_shaping_params),
             layout_name=self.layout_name,
-            start_all_orders=self.start_all_orders
+            start_all_orders=self.start_all_orders,
         )
 
     @property
@@ -990,10 +1253,9 @@ def mdp_params(self):
             "start_player_positions": self.start_player_positions,
             "start_bonus_orders": self.start_bonus_orders,
             "rew_shaping_params": copy.deepcopy(self.reward_shaping_params),
-            "start_all_orders" : self.start_all_orders
+            "start_all_orders": self.start_all_orders,
         }
 
-
     ##############
     # GAME LOGIC #
     ##############
@@ -1006,34 +1268,49 @@ def get_actions(self, state):
         can take.
         """
         self._check_valid_state(state)
-        return [self._get_player_actions(state, i) for i in range(len(state.players))]
+        return [
+            self._get_player_actions(state, i)
+            for i in range(len(state.players))
+        ]
 
     def _get_player_actions(self, state, player_num):
         """All actions are allowed to all players in all states."""
         return Action.ALL_ACTIONS
 
     def _check_action(self, state, joint_action):
-        for p_action, p_legal_actions in zip(joint_action, self.get_actions(state)):
+        for p_action, p_legal_actions in zip(
+            joint_action, self.get_actions(state)
+        ):
             if p_action not in p_legal_actions:
-                raise ValueError('Invalid action')
+                raise ValueError("Invalid action")
 
     def get_standard_start_state(self):
         if self.start_state:
             return self.start_state
         start_state = OvercookedState.from_player_positions(
-            self.start_player_positions, bonus_orders=self.start_bonus_orders, all_orders=self.start_all_orders
+            self.start_player_positions,
+            bonus_orders=self.start_bonus_orders,
+            all_orders=self.start_all_orders,
         )
         return start_state
 
-    def get_random_start_state_fn(self, random_start_pos=False, rnd_obj_prob_thresh=0.0):
+    def get_random_start_state_fn(
+        self, random_start_pos=False, rnd_obj_prob_thresh=0.0
+    ):
         def start_state_fn():
             if random_start_pos:
                 valid_positions = self.get_valid_joint_player_positions()
-                start_pos = valid_positions[np.random.choice(len(valid_positions))]
+                start_pos = valid_positions[
+                    np.random.choice(len(valid_positions))
+                ]
             else:
                 start_pos = self.start_player_positions
 
-            start_state = OvercookedState.from_player_positions(start_pos, bonus_orders=self.start_bonus_orders, all_orders=self.start_all_orders)
+            start_state = OvercookedState.from_player_positions(
+                start_pos,
+                bonus_orders=self.start_bonus_orders,
+                all_orders=self.start_all_orders,
+            )
 
             if rnd_obj_prob_thresh == 0:
                 return start_state
@@ -1046,56 +1323,80 @@ def start_state_fn():
                 p = np.random.rand()
                 if p < rnd_obj_prob_thresh:
                     n = int(np.random.randint(low=1, high=4))
-                    m = int(np.random.randint(low=0, high=4-n))
+                    m = int(np.random.randint(low=0, high=4 - n))
                     q = np.random.rand()
                     cooking_tick = 0 if q < rnd_obj_prob_thresh else -1
-                    start_state.objects[pot_loc] = SoupState.get_soup(pot_loc, num_onions=n, num_tomatoes=m, cooking_tick=cooking_tick)
+                    start_state.objects[pot_loc] = SoupState.get_soup(
+                        pot_loc,
+                        num_onions=n,
+                        num_tomatoes=m,
+                        cooking_tick=cooking_tick,
+                    )
 
             # For each player, add a random object with prob rnd_obj_prob_thresh
             for player in start_state.players:
                 p = np.random.rand()
                 if p < rnd_obj_prob_thresh:
                     # Different objects have different probabilities
-                    obj = np.random.choice(["dish", "onion", "soup"], p=[0.2, 0.6, 0.2])
+                    obj = np.random.choice(
+                        ["dish", "onion", "soup"], p=[0.2, 0.6, 0.2]
+                    )
                     n = int(np.random.randint(low=1, high=4))
-                    m = int(np.random.randint(low=0, high=4-n))
+                    m = int(np.random.randint(low=0, high=4 - n))
                     if obj == "soup":
                         player.set_object(
-                            SoupState.get_soup(player.position, num_onions=n, num_tomatoes=m, finished=True)
+                            SoupState.get_soup(
+                                player.position,
+                                num_onions=n,
+                                num_tomatoes=m,
+                                finished=True,
+                            )
                         )
                     else:
                         player.set_object(ObjectState(obj, player.position))
             return start_state
+
         return start_state_fn
 
     def is_terminal(self, state):
         # There is a finite horizon, handled by the environment.
         return False
 
-    def get_state_transition(self, state, joint_action, display_phi=False, motion_planner=None):
+    def get_state_transition(
+        self, state, joint_action, display_phi=False, motion_planner=None
+    ):
         """Gets information about possible transitions for the action.
 
         Returns the next state, sparse reward and reward shaping.
         Assumes all actions are deterministic.
 
-        NOTE: Sparse reward is given only when soups are delivered, 
-        shaped reward is given only for completion of subgoals 
+        NOTE: Sparse reward is given only when soups are delivered,
+        shaped reward is given only for completion of subgoals
         (not soup deliveries).
         """
-        events_infos = { event : [False] * self.num_players for event in EVENT_TYPES }
-        assert not self.is_terminal(state), "Trying to find successor of a terminal state: {}".format(state)
+        events_infos = {
+            event: [False] * self.num_players for event in EVENT_TYPES
+        }
+        assert not self.is_terminal(
+            state
+        ), "Trying to find successor of a terminal state: {}".format(state)
         for action, action_set in zip(joint_action, self.get_actions(state)):
             if action not in action_set:
-                raise ValueError("Illegal action %s in state %s" % (action, state))
-        
+                raise ValueError(
+                    "Illegal action %s in state %s" % (action, state)
+                )
+
         new_state = state.deepcopy()
 
         # Resolve interacts first
-        sparse_reward_by_agent, shaped_reward_by_agent = self.resolve_interacts(new_state, joint_action, events_infos)
+        (
+            sparse_reward_by_agent,
+            shaped_reward_by_agent,
+        ) = self.resolve_interacts(new_state, joint_action, events_infos)
 
         assert new_state.player_positions == state.player_positions
         assert new_state.player_orientations == state.player_orientations
-        
+
         # Resolve player movements
         self.resolve_movement(new_state, joint_action)
 
@@ -1110,23 +1411,31 @@ def get_state_transition(self, state, joint_action, display_phi=False, motion_pl
             "shaped_reward_by_agent": shaped_reward_by_agent,
         }
         if display_phi:
-            assert motion_planner is not None, "motion planner must be defined if display_phi is true"
+            assert (
+                motion_planner is not None
+            ), "motion planner must be defined if display_phi is true"
             infos["phi_s"] = self.potential_function(state, motion_planner)
-            infos["phi_s_prime"] = self.potential_function(new_state, motion_planner)
+            infos["phi_s_prime"] = self.potential_function(
+                new_state, motion_planner
+            )
         return new_state, infos
 
     def resolve_interacts(self, new_state, joint_action, events_infos):
         """
         Resolve any INTERACT actions, if present.
 
-        Currently if two players both interact with a terrain, we resolve player 1's interact 
+        Currently if two players both interact with a terrain, we resolve player 1's interact
         first and then player 2's, without doing anything like collision checking.
         """
         pot_states = self.get_pot_states(new_state)
         # We divide reward by agent to keep track of who contributed
-        sparse_reward, shaped_reward = [0] * self.num_players, [0] * self.num_players 
+        sparse_reward, shaped_reward = [0] * self.num_players, [
+            0
+        ] * self.num_players
 
-        for player_idx, (player, action) in enumerate(zip(new_state.players, joint_action)):
+        for player_idx, (player, action) in enumerate(
+            zip(new_state.players, joint_action)
+        ):
 
             if action != Action.INTERACT:
                 continue
@@ -1137,63 +1446,87 @@ def resolve_interacts(self, new_state, joint_action, events_infos):
 
             # NOTE: we always log pickup/drop before performing it, as that's
             # what the logic of determining whether the pickup/drop is useful assumes
-            if terrain_type == 'X':
+            if terrain_type == "X":
 
                 if player.has_object() and not new_state.has_object(i_pos):
                     obj_name = player.get_object().name
-                    self.log_object_drop(events_infos, new_state, obj_name, pot_states, player_idx)
+                    self.log_object_drop(
+                        events_infos,
+                        new_state,
+                        obj_name,
+                        pot_states,
+                        player_idx,
+                    )
 
                     # Drop object on counter
                     obj = player.remove_object()
                     new_state.add_object(obj, i_pos)
-                    
+
                 elif not player.has_object() and new_state.has_object(i_pos):
                     obj_name = new_state.get_object(i_pos).name
-                    self.log_object_pickup(events_infos, new_state, obj_name, pot_states, player_idx)
+                    self.log_object_pickup(
+                        events_infos,
+                        new_state,
+                        obj_name,
+                        pot_states,
+                        player_idx,
+                    )
 
                     # Pick up object from counter
                     obj = new_state.remove_object(i_pos)
                     player.set_object(obj)
-                    
 
-            elif terrain_type == 'O' and player.held_object is None:
-                self.log_object_pickup(events_infos, new_state, "onion", pot_states, player_idx)
+            elif terrain_type == "O" and player.held_object is None:
+                self.log_object_pickup(
+                    events_infos, new_state, "onion", pot_states, player_idx
+                )
 
                 # Onion pickup from dispenser
-                obj = ObjectState('onion', pos)
+                obj = ObjectState("onion", pos)
                 player.set_object(obj)
 
-            elif terrain_type == 'T' and player.held_object is None:
+            elif terrain_type == "T" and player.held_object is None:
                 # Tomato pickup from dispenser
-                player.set_object(ObjectState('tomato', pos))
+                player.set_object(ObjectState("tomato", pos))
 
-            elif terrain_type == 'D' and player.held_object is None:
-                self.log_object_pickup(events_infos, new_state, "dish", pot_states, player_idx)
+            elif terrain_type == "D" and player.held_object is None:
+                self.log_object_pickup(
+                    events_infos, new_state, "dish", pot_states, player_idx
+                )
 
                 # Give shaped reward if pickup is useful
                 if self.is_dish_pickup_useful(new_state, pot_states):
-                    shaped_reward[player_idx] += self.reward_shaping_params["DISH_PICKUP_REWARD"]
+                    shaped_reward[player_idx] += self.reward_shaping_params[
+                        "DISH_PICKUP_REWARD"
+                    ]
 
                 # Perform dish pickup from dispenser
-                obj = ObjectState('dish', pos)
+                obj = ObjectState("dish", pos)
                 player.set_object(obj)
 
-            elif terrain_type == 'P' and not player.has_object():
+            elif terrain_type == "P" and not player.has_object():
                 # Cooking soup
                 if self.soup_to_be_cooked_at_location(new_state, i_pos):
                     soup = new_state.get_object(i_pos)
                     soup.begin_cooking()
-            
-            elif terrain_type == 'P' and player.has_object():
 
-                if player.get_object().name == 'dish' and self.soup_ready_at_location(new_state, i_pos):
-                    self.log_object_pickup(events_infos, new_state, "soup", pot_states, player_idx)
+            elif terrain_type == "P" and player.has_object():
+
+                if (
+                    player.get_object().name == "dish"
+                    and self.soup_ready_at_location(new_state, i_pos)
+                ):
+                    self.log_object_pickup(
+                        events_infos, new_state, "soup", pot_states, player_idx
+                    )
 
                     # Pick up soup
-                    player.remove_object() # Remove the dish
-                    obj = new_state.remove_object(i_pos) # Get soup
+                    player.remove_object()  # Remove the dish
+                    obj = new_state.remove_object(i_pos)  # Get soup
                     player.set_object(obj)
-                    shaped_reward[player_idx] += self.reward_shaping_params["SOUP_PICKUP_REWARD"]
+                    shaped_reward[player_idx] += self.reward_shaping_params[
+                        "SOUP_PICKUP_REWARD"
+                    ]
 
                 elif player.get_object().name in Recipe.ALL_INGREDIENTS:
                     # Adding ingredient to soup
@@ -1208,26 +1541,42 @@ def resolve_interacts(self, new_state, joint_action, events_infos):
                         old_soup = soup.deepcopy()
                         obj = player.remove_object()
                         soup.add_ingredient(obj)
-                        shaped_reward[player_idx] += self.reward_shaping_params["PLACEMENT_IN_POT_REW"]
+                        shaped_reward[
+                            player_idx
+                        ] += self.reward_shaping_params["PLACEMENT_IN_POT_REW"]
 
                         # Log potting
-                        self.log_object_potting(events_infos, new_state, old_soup, soup, obj.name, player_idx)
+                        self.log_object_potting(
+                            events_infos,
+                            new_state,
+                            old_soup,
+                            soup,
+                            obj.name,
+                            player_idx,
+                        )
                         if obj.name == Recipe.ONION:
-                            events_infos['potting_onion'][player_idx] = True
+                            events_infos["potting_onion"][player_idx] = True
 
-            elif terrain_type == 'S' and player.has_object():
+            elif terrain_type == "S" and player.has_object():
                 obj = player.get_object()
-                if obj.name == 'soup':
+                if obj.name == "soup":
 
                     delivery_rew = self.deliver_soup(new_state, player, obj)
                     sparse_reward[player_idx] += delivery_rew
 
                     # Log soup delivery
-                    events_infos['soup_delivery'][player_idx] = True                        
+                    events_infos["soup_delivery"][player_idx] = True
 
         return sparse_reward, shaped_reward
 
-    def get_recipe_value(self, state, recipe, discounted=False, base_recipe=None, potential_params={}):
+    def get_recipe_value(
+        self,
+        state,
+        recipe,
+        discounted=False,
+        base_recipe=None,
+        potential_params={},
+    ):
         """
         Return the reward the player should receive for delivering this recipe
 
@@ -1237,7 +1586,7 @@ def get_recipe_value(self, state, recipe, discounted=False, base_recipe=None, po
         if not discounted:
             if not recipe in state.all_orders:
                 return 0
-            
+
             if not recipe in state.bonus_orders:
                 return recipe.value
 
@@ -1245,22 +1594,39 @@ def get_recipe_value(self, state, recipe, discounted=False, base_recipe=None, po
         else:
             # Calculate missing ingredients needed to complete recipe
             missing_ingredients = list(recipe.ingredients)
-            prev_ingredients = list(base_recipe.ingredients) if base_recipe else []
+            prev_ingredients = (
+                list(base_recipe.ingredients) if base_recipe else []
+            )
             for ingredient in prev_ingredients:
                 missing_ingredients.remove(ingredient)
-            n_tomatoes = len([i for i in missing_ingredients if i == Recipe.TOMATO])
-            n_onions = len([i for i in missing_ingredients if i == Recipe.ONION])
+            n_tomatoes = len(
+                [i for i in missing_ingredients if i == Recipe.TOMATO]
+            )
+            n_onions = len(
+                [i for i in missing_ingredients if i == Recipe.ONION]
+            )
 
-            gamma, pot_onion_steps, pot_tomato_steps = potential_params['gamma'], potential_params['pot_onion_steps'], potential_params['pot_tomato_steps']
+            gamma, pot_onion_steps, pot_tomato_steps = (
+                potential_params["gamma"],
+                potential_params["pot_onion_steps"],
+                potential_params["pot_tomato_steps"],
+            )
 
-            return gamma**recipe.time * gamma**(pot_onion_steps * n_onions) * gamma**(pot_tomato_steps * n_tomatoes) * self.get_recipe_value(state, recipe, discounted=False)
+            return (
+                gamma**recipe.time
+                * gamma ** (pot_onion_steps * n_onions)
+                * gamma ** (pot_tomato_steps * n_tomatoes)
+                * self.get_recipe_value(state, recipe, discounted=False)
+            )
 
     def deliver_soup(self, state, player, soup):
         """
         Deliver the soup, and get reward if there is no order list
         or if the type of the delivered soup matches the next order.
         """
-        assert soup.name == 'soup', "Tried to deliver something that wasn't soup"
+        assert (
+            soup.name == "soup"
+        ), "Tried to deliver something that wasn't soup"
         assert soup.is_ready, "Tried to deliever soup that isn't ready"
         player.remove_object()
 
@@ -1268,15 +1634,29 @@ def deliver_soup(self, state, player, soup):
 
     def resolve_movement(self, state, joint_action):
         """Resolve player movement and deal with possible collisions"""
-        new_positions, new_orientations = self.compute_new_positions_and_orientations(state.players, joint_action)
-        for player_state, new_pos, new_o in zip(state.players, new_positions, new_orientations):
+        (
+            new_positions,
+            new_orientations,
+        ) = self.compute_new_positions_and_orientations(
+            state.players, joint_action
+        )
+        for player_state, new_pos, new_o in zip(
+            state.players, new_positions, new_orientations
+        ):
             player_state.update_pos_and_or(new_pos, new_o)
 
-    def compute_new_positions_and_orientations(self, old_player_states, joint_action):
+    def compute_new_positions_and_orientations(
+        self, old_player_states, joint_action
+    ):
         """Compute new positions and orientations ignoring collisions"""
-        new_positions, new_orientations = list(zip(*[
-            self._move_if_direction(p.position, p.orientation, a) \
-            for p, a in zip(old_player_states, joint_action)]))
+        new_positions, new_orientations = list(
+            zip(
+                *[
+                    self._move_if_direction(p.position, p.orientation, a)
+                    for p, a in zip(old_player_states, joint_action)
+                ]
+            )
+        )
         old_positions = tuple(p.position for p in old_player_states)
         new_positions = self._handle_collisions(old_positions, new_positions)
         return new_positions, new_orientations
@@ -1294,12 +1674,15 @@ def is_transition_collision(self, old_positions, new_positions):
         return False
 
     def is_joint_position_collision(self, joint_position):
-        return any(pos0 == pos1 for pos0, pos1 in itertools.combinations(joint_position, 2))
-            
+        return any(
+            pos0 == pos1
+            for pos0, pos1 in itertools.combinations(joint_position, 2)
+        )
+
     def step_environment_effects(self, state):
         state.timestep += 1
         for obj in state.objects.values():
-            if obj.name == 'soup' and obj.is_cooking:
+            if obj.name == "soup" and obj.is_cooking:
                 obj.cook()
 
     def _handle_collisions(self, old_positions, new_positions):
@@ -1316,7 +1699,7 @@ def _get_terrain_type_pos_dict(self):
         return pos_dict
 
     def _move_if_direction(self, position, orientation, action):
-        """Returns position and orientation that would 
+        """Returns position and orientation that would
         be obtained after executing action"""
         if action not in Action.MOTION_ACTIONS:
             return position, orientation
@@ -1326,19 +1709,24 @@ def _move_if_direction(self, position, orientation, action):
             return position, new_orientation
         return new_pos, new_orientation
 
-
     #######################
     # LAYOUT / STATE INFO #
     #######################
 
     def get_valid_player_positions(self):
-        return self.terrain_pos_dict[' ']
+        return self.terrain_pos_dict[" "]
 
     def get_valid_joint_player_positions(self):
         """Returns all valid tuples of the form (p0_pos, p1_pos, p2_pos, ...)"""
-        valid_positions = self.get_valid_player_positions() 
-        all_joint_positions = list(itertools.product(valid_positions, repeat=self.num_players))
-        valid_joint_positions = [j_pos for j_pos in all_joint_positions if not self.is_joint_position_collision(j_pos)]
+        valid_positions = self.get_valid_player_positions()
+        all_joint_positions = list(
+            itertools.product(valid_positions, repeat=self.num_players)
+        )
+        valid_joint_positions = [
+            j_pos
+            for j_pos in all_joint_positions
+            if not self.is_joint_position_collision(j_pos)
+        ]
         return valid_joint_positions
 
     def get_valid_player_positions_and_orientations(self):
@@ -1350,11 +1738,18 @@ def get_valid_player_positions_and_orientations(self):
     def get_valid_joint_player_positions_and_orientations(self):
         """All joint player position and orientation pairs that are not
         overlapping and on empty terrain."""
-        valid_player_states = self.get_valid_player_positions_and_orientations()
+        valid_player_states = (
+            self.get_valid_player_positions_and_orientations()
+        )
 
         valid_joint_player_states = []
-        for players_pos_and_orientations in itertools.product(valid_player_states, repeat=self.num_players):
-            joint_position = [plyer_pos_and_or[0] for plyer_pos_and_or in players_pos_and_orientations]
+        for players_pos_and_orientations in itertools.product(
+            valid_player_states, repeat=self.num_players
+        ):
+            joint_position = [
+                plyer_pos_and_or[0]
+                for plyer_pos_and_or in players_pos_and_orientations
+            ]
             if not self.is_joint_position_collision(joint_position):
                 valid_joint_player_states.append(players_pos_and_orientations)
 
@@ -1373,22 +1768,22 @@ def get_terrain_type_at_pos(self, pos):
         return self.terrain_mtx[y][x]
 
     def get_dish_dispenser_locations(self):
-        return list(self.terrain_pos_dict['D'])
+        return list(self.terrain_pos_dict["D"])
 
     def get_onion_dispenser_locations(self):
-        return list(self.terrain_pos_dict['O'])
+        return list(self.terrain_pos_dict["O"])
 
     def get_tomato_dispenser_locations(self):
-        return list(self.terrain_pos_dict['T'])
+        return list(self.terrain_pos_dict["T"])
 
     def get_serving_locations(self):
-        return list(self.terrain_pos_dict['S'])
+        return list(self.terrain_pos_dict["S"])
 
     def get_pot_locations(self):
-        return list(self.terrain_pos_dict['P'])
+        return list(self.terrain_pos_dict["P"])
 
     def get_counter_locations(self):
-        return list(self.terrain_pos_dict['X'])
+        return list(self.terrain_pos_dict["X"])
 
     @property
     def num_pots(self):
@@ -1407,23 +1802,31 @@ def get_pot_states(self, state):
         pots_states_dict = defaultdict(list)
         for pot_pos in self.get_pot_locations():
             if not state.has_object(pot_pos):
-                pots_states_dict['empty'].append(pot_pos)
+                pots_states_dict["empty"].append(pot_pos)
             else:
                 soup = state.get_object(pot_pos)
-                assert soup.name == 'soup', "soup at " + pot_pos + " is not a soup but a " + soup.name
+                assert soup.name == "soup", (
+                    "soup at " + pot_pos + " is not a soup but a " + soup.name
+                )
                 if soup.is_ready:
-                    pots_states_dict['ready'].append(pot_pos)
+                    pots_states_dict["ready"].append(pot_pos)
                 elif soup.is_cooking:
-                    pots_states_dict['cooking'].append(pot_pos)
+                    pots_states_dict["cooking"].append(pot_pos)
                 else:
                     num_ingredients = len(soup.ingredients)
-                    pots_states_dict['{}_items'.format(num_ingredients)].append(pot_pos)
+                    pots_states_dict[
+                        "{}_items".format(num_ingredients)
+                    ].append(pot_pos)
 
         return pots_states_dict
 
     def get_counter_objects_dict(self, state, counter_subset=None):
         """Returns a dictionary of pos:objects on counters by type"""
-        counters_considered = self.terrain_pos_dict['X'] if counter_subset is None else counter_subset
+        counters_considered = (
+            self.terrain_pos_dict["X"]
+            if counter_subset is None
+            else counter_subset
+        )
         counter_objects_dict = defaultdict(list)
         for obj in state.objects.values():
             if obj.position in counters_considered:
@@ -1439,35 +1842,53 @@ def get_empty_pots(self, pot_states):
         return pot_states["empty"]
 
     def get_non_empty_pots(self, pot_states):
-        return self.get_full_pots(pot_states) + self.get_partially_full_pots(pot_states)
+        return self.get_full_pots(pot_states) + self.get_partially_full_pots(
+            pot_states
+        )
 
     def get_ready_pots(self, pot_states):
-        return pot_states['ready']
+        return pot_states["ready"]
 
     def get_cooking_pots(self, pot_states):
-        return pot_states['cooking']
+        return pot_states["cooking"]
 
     def get_full_but_not_cooking_pots(self, pot_states):
-        return pot_states['{}_items'.format(Recipe.MAX_NUM_INGREDIENTS)]
+        return pot_states["{}_items".format(Recipe.MAX_NUM_INGREDIENTS)]
 
     def get_full_pots(self, pot_states):
-        return self.get_cooking_pots(pot_states) + self.get_ready_pots(pot_states) + self.get_full_but_not_cooking_pots(pot_states)
+        return (
+            self.get_cooking_pots(pot_states)
+            + self.get_ready_pots(pot_states)
+            + self.get_full_but_not_cooking_pots(pot_states)
+        )
 
     def get_partially_full_pots(self, pot_states):
-        return list(set().union(*[pot_states['{}_items'.format(i)] for i in range(1, Recipe.MAX_NUM_INGREDIENTS)]))
+        return list(
+            set().union(
+                *[
+                    pot_states["{}_items".format(i)]
+                    for i in range(1, Recipe.MAX_NUM_INGREDIENTS)
+                ]
+            )
+        )
 
     def soup_ready_at_location(self, state, pos):
         if not state.has_object(pos):
             return False
         obj = state.get_object(pos)
-        assert obj.name == 'soup', 'Object in pot was not soup'
+        assert obj.name == "soup", "Object in pot was not soup"
         return obj.is_ready
 
     def soup_to_be_cooked_at_location(self, state, pos):
         if not state.has_object(pos):
             return False
         obj = state.get_object(pos)
-        return obj.name == 'soup' and not obj.is_cooking and not obj.is_ready and len(obj.ingredients) > 0
+        return (
+            obj.name == "soup"
+            and not obj.is_cooking
+            and not obj.is_ready
+            and len(obj.ingredients) > 0
+        )
 
     def _check_valid_state(self, state):
         """Checks that the state is valid.
@@ -1488,19 +1909,23 @@ def _check_valid_state(self, state):
             # Check that held objects have the same position
             if player_state.held_object is not None:
                 all_objects.append(player_state.held_object)
-                assert player_state.held_object.position == player_state.position
+                assert (
+                    player_state.held_object.position == player_state.position
+                )
 
         for obj_pos, obj_state in state.objects.items():
             # Check that the hash key position agrees with the position stored
             # in the object state
             assert obj_state.position == obj_pos
             # Check that non-held objects are on terrain
-            assert self.get_terrain_type_at_pos(obj_pos) != ' '
+            assert self.get_terrain_type_at_pos(obj_pos) != " "
 
         # Check that players and non-held objects don't overlap
         all_pos = [player_state.position for player_state in state.players]
         all_pos += [obj_state.position for obj_state in state.objects.values()]
-        assert len(all_pos) == len(set(all_pos)), "Overlapping players or objects"
+        assert len(all_pos) == len(
+            set(all_pos)
+        ), "Overlapping players or objects"
 
         # Check that objects have a valid state
         for obj_state in all_objects:
@@ -1513,12 +1938,27 @@ def find_free_counters_valid_for_both_players(self, state, mlam):
         free_counters_valid_for_both = []
         for free_counter in free_counters:
             goals = mlam.motion_planner.motion_goals_for_pos[free_counter]
-            if any([mlam.motion_planner.is_valid_motion_start_goal_pair(one_player.pos_and_or, goal) for goal in goals]) and \
-            any([mlam.motion_planner.is_valid_motion_start_goal_pair(other_player.pos_and_or, goal) for goal in goals]):
+            if any(
+                [
+                    mlam.motion_planner.is_valid_motion_start_goal_pair(
+                        one_player.pos_and_or, goal
+                    )
+                    for goal in goals
+                ]
+            ) and any(
+                [
+                    mlam.motion_planner.is_valid_motion_start_goal_pair(
+                        other_player.pos_and_or, goal
+                    )
+                    for goal in goals
+                ]
+            ):
                 free_counters_valid_for_both.append(free_counter)
         return free_counters_valid_for_both
 
-    def _get_optimal_possible_recipe(self, state, recipe, discounted, potential_params, return_value):
+    def _get_optimal_possible_recipe(
+        self, state, recipe, discounted, potential_params, return_value
+    ):
         """
         Traverse the recipe-space graph using DFS to find the best possible recipe that can be made
         from the current recipe
@@ -1540,27 +1980,41 @@ def _get_optimal_possible_recipe(self, state, recipe, discounted, potential_para
             curr_recipe = stack.pop()
             if curr_recipe not in visited:
                 visited.add(curr_recipe)
-                curr_value = self.get_recipe_value(state, curr_recipe, base_recipe=start_recipe, discounted=discounted, potential_params=potential_params)
+                curr_value = self.get_recipe_value(
+                    state,
+                    curr_recipe,
+                    base_recipe=start_recipe,
+                    discounted=discounted,
+                    potential_params=potential_params,
+                )
                 if curr_value > best_value:
                     best_value, best_recipe = curr_value, curr_recipe
-                
+
                 for neighbor in curr_recipe.neighbors():
                     if not neighbor in visited:
                         stack.append(neighbor)
-        
+
         if return_value:
             return best_recipe, best_value
         return best_recipe
 
-
-    def get_optimal_possible_recipe(self, state, recipe, discounted=False, potential_params={}, return_value=False):
+    def get_optimal_possible_recipe(
+        self,
+        state,
+        recipe,
+        discounted=False,
+        potential_params={},
+        return_value=False,
+    ):
         """
         Return the best possible recipe that can be made starting with ingredients in `recipe`
         Uses self._optimal_possible_recipe as a cache to avoid re-computing. This only works because
         the recipe values are currently static (i.e. bonus_orders doesn't change). Would need to have cache
         flushed if order dynamics are introduced
         """
-        cache_valid = not discounted or self._prev_potential_params == potential_params
+        cache_valid = (
+            not discounted or self._prev_potential_params == potential_params
+        )
         if not cache_valid:
             if discounted:
                 self._opt_recipe_discount_cache = {}
@@ -1575,16 +2029,19 @@ def get_optimal_possible_recipe(self, state, recipe, discounted=False, potential
 
         if recipe not in cache:
             # Compute best recipe now and store in cache for later use
-            opt_recipe, value = self._get_optimal_possible_recipe(state, recipe, discounted=discounted, potential_params=potential_params, return_value=True)
+            opt_recipe, value = self._get_optimal_possible_recipe(
+                state,
+                recipe,
+                discounted=discounted,
+                potential_params=potential_params,
+                return_value=True,
+            )
             cache[recipe] = (opt_recipe, value)
 
         # Return best recipe (and value) from cache
         if return_value:
             return cache[recipe]
         return cache[recipe][0]
-        
-
-
 
     @staticmethod
     def _assert_valid_grid(grid):
@@ -1600,40 +2057,53 @@ def _assert_valid_grid(grid):
         width = len(grid[0])
 
         # Make sure the grid is not ragged
-        assert all(len(row) == width for row in grid), 'Ragged grid'
+        assert all(len(row) == width for row in grid), "Ragged grid"
 
         # Borders must not be free spaces
         def is_not_free(c):
-            return c in 'XOPDST'
+            return c in "XOPDST"
 
         for y in range(height):
-            assert is_not_free(grid[y][0]), 'Left border must not be free'
-            assert is_not_free(grid[y][-1]), 'Right border must not be free'
+            assert is_not_free(grid[y][0]), "Left border must not be free"
+            assert is_not_free(grid[y][-1]), "Right border must not be free"
         for x in range(width):
-            assert is_not_free(grid[0][x]), 'Top border must not be free'
-            assert is_not_free(grid[-1][x]), 'Bottom border must not be free'
+            assert is_not_free(grid[0][x]), "Top border must not be free"
+            assert is_not_free(grid[-1][x]), "Bottom border must not be free"
 
         all_elements = [element for row in grid for element in row]
-        digits = ['1', '2', '3', '4', '5', '6', '7', '8', '9']
+        digits = ["1", "2", "3", "4", "5", "6", "7", "8", "9"]
         layout_digits = [e for e in all_elements if e in digits]
         num_players = len(layout_digits)
         assert num_players > 0, "No players (digits) in grid"
         layout_digits = list(sorted(map(int, layout_digits)))
-        assert layout_digits == list(range(1, num_players + 1)), "Some players were missing"
-
-        assert all(c in 'XOPDST123456789 ' for c in all_elements), 'Invalid character in grid'
-        assert all_elements.count('1') == 1, "'1' must be present exactly once"
-        assert all_elements.count('D') >= 1, "'D' must be present at least once"
-        assert all_elements.count('S') >= 1, "'S' must be present at least once"
-        assert all_elements.count('P') >= 1, "'P' must be present at least once"
-        assert all_elements.count('O') >= 1 or all_elements.count('T') >= 1, "'O' or 'T' must be present at least once"
-
+        assert layout_digits == list(
+            range(1, num_players + 1)
+        ), "Some players were missing"
+
+        assert all(
+            c in "XOPDST123456789 " for c in all_elements
+        ), "Invalid character in grid"
+        assert all_elements.count("1") == 1, "'1' must be present exactly once"
+        assert (
+            all_elements.count("D") >= 1
+        ), "'D' must be present at least once"
+        assert (
+            all_elements.count("S") >= 1
+        ), "'S' must be present at least once"
+        assert (
+            all_elements.count("P") >= 1
+        ), "'P' must be present at least once"
+        assert (
+            all_elements.count("O") >= 1 or all_elements.count("T") >= 1
+        ), "'O' or 'T' must be present at least once"
 
     ################################
     # EVENT LOGGING HELPER METHODS #
     ################################
 
-    def log_object_potting(self, events_infos, state, old_soup, new_soup, obj_name, player_index):
+    def log_object_potting(
+        self, events_infos, state, old_soup, new_soup, obj_name, player_index
+    ):
         """Player added an ingredient to a pot"""
         obj_pickup_key = "potting_" + obj_name
         if obj_pickup_key not in events_infos:
@@ -1641,10 +2111,10 @@ def log_object_potting(self, events_infos, state, old_soup, new_soup, obj_name,
         events_infos[obj_pickup_key][player_index] = True
 
         POTTING_FNS = {
-            "optimal" : self.is_potting_optimal,
-            "catastrophic" : self.is_potting_catastrophic,
-            "viable" : self.is_potting_viable,
-            "useless" : self.is_potting_useless
+            "optimal": self.is_potting_optimal,
+            "catastrophic": self.is_potting_catastrophic,
+            "viable": self.is_potting_viable,
+            "useless": self.is_potting_useless,
         }
 
         for outcome, outcome_fn in POTTING_FNS.items():
@@ -1652,35 +2122,38 @@ def log_object_potting(self, events_infos, state, old_soup, new_soup, obj_name,
                 potting_key = "{}_{}_potting".format(outcome, obj_name)
                 events_infos[potting_key][player_index] = True
 
-    
-    def log_object_pickup(self, events_infos, state, obj_name, pot_states, player_index):
+    def log_object_pickup(
+        self, events_infos, state, obj_name, pot_states, player_index
+    ):
         """Player picked an object up from a counter or a dispenser"""
         obj_pickup_key = obj_name + "_pickup"
         if obj_pickup_key not in events_infos:
             raise ValueError("Unknown event {}".format(obj_pickup_key))
         events_infos[obj_pickup_key][player_index] = True
-        
+
         USEFUL_PICKUP_FNS = {
-            "tomato" : self.is_ingredient_pickup_useful,
+            "tomato": self.is_ingredient_pickup_useful,
             "onion": self.is_ingredient_pickup_useful,
-            "dish": self.is_dish_pickup_useful
+            "dish": self.is_dish_pickup_useful,
         }
         if obj_name in USEFUL_PICKUP_FNS:
             if USEFUL_PICKUP_FNS[obj_name](state, pot_states, player_index):
                 obj_useful_key = "useful_" + obj_name + "_pickup"
                 events_infos[obj_useful_key][player_index] = True
 
-    def log_object_drop(self, events_infos, state, obj_name, pot_states, player_index):
+    def log_object_drop(
+        self, events_infos, state, obj_name, pot_states, player_index
+    ):
         """Player dropped the object on a counter"""
         obj_drop_key = obj_name + "_drop"
         if obj_drop_key not in events_infos:
             raise ValueError("Unknown event {}".format(obj_drop_key))
         events_infos[obj_drop_key][player_index] = True
-        
+
         USEFUL_DROP_FNS = {
-            "tomato" : self.is_ingredient_drop_useful,
+            "tomato": self.is_ingredient_drop_useful,
             "onion": self.is_ingredient_drop_useful,
-            "dish": self.is_dish_drop_useful
+            "dish": self.is_dish_drop_useful,
         }
         if obj_name in USEFUL_DROP_FNS:
             if USEFUL_DROP_FNS[obj_name](state, pot_states, player_index):
@@ -1691,21 +2164,26 @@ def is_dish_pickup_useful(self, state, pot_states, player_index=None):
         """
         NOTE: this only works if self.num_players == 2
         Useful if:
-        - Pot is ready/cooking and there is no player with a dish               \ 
+        - Pot is ready/cooking and there is no player with a dish               \
         - 2 pots are ready/cooking and there is one player with a dish          | -> number of dishes in players hands < number of ready/cooking/partially full soups 
         - Partially full pot is ok if the other player is on course to fill it  /
 
         We also want to prevent picking up and dropping dishes, so add the condition
         that there must be no dishes on counters
         """
-        if self.num_players != 2: return False
+        if self.num_players != 2:
+            return False
 
         # This next line is to prevent reward hacking (this logic is also used by reward shaping)
         dishes_on_counters = self.get_counter_objects_dict(state)["dish"]
         no_dishes_on_counters = len(dishes_on_counters) == 0
 
-        num_player_dishes = len(state.player_objects_by_type['dish'])
-        non_empty_pots = len(self.get_ready_pots(pot_states) + self.get_cooking_pots(pot_states) + self.get_partially_full_pots(pot_states))
+        num_player_dishes = len(state.player_objects_by_type["dish"])
+        non_empty_pots = len(
+            self.get_ready_pots(pot_states)
+            + self.get_cooking_pots(pot_states)
+            + self.get_partially_full_pots(pot_states)
+        )
         return no_dishes_on_counters and num_player_dishes < non_empty_pots
 
     def is_dish_drop_useful(self, state, pot_states, player_index):
@@ -1715,10 +2193,14 @@ def is_dish_drop_useful(self, state, pot_states, player_index):
         - Onion is needed (all pots are non-full)
         - Nobody is holding onions
         """
-        if self.num_players != 2: return False
+        if self.num_players != 2:
+            return False
         all_non_full = len(self.get_full_pots(pot_states)) == 0
         other_player = state.players[1 - player_index]
-        other_player_holding_onion = other_player.has_object() and other_player.get_object().name == "onion"
+        other_player_holding_onion = (
+            other_player.has_object()
+            and other_player.get_object().name == "onion"
+        )
         return all_non_full and not other_player_holding_onion
 
     def is_ingredient_pickup_useful(self, state, pot_states, player_index):
@@ -1727,12 +2209,16 @@ def is_ingredient_pickup_useful(self, state, pot_states, player_index):
         Always useful unless:
         - All pots are full & other agent is not holding a dish
         """
-        if self.num_players != 2: return False
+        if self.num_players != 2:
+            return False
         all_pots_full = self.num_pots == len(self.get_full_pots(pot_states))
         other_player = state.players[1 - player_index]
-        other_player_has_dish = other_player.has_object() and other_player.get_object().name == "dish"
+        other_player_has_dish = (
+            other_player.has_object()
+            and other_player.get_object().name == "dish"
+        )
         return not (all_pots_full and not other_player_has_dish)
-        
+
     def is_ingredient_drop_useful(self, state, pot_states, player_index):
         """
         NOTE: this only works if self.num_players == 2
@@ -1740,20 +2226,30 @@ def is_ingredient_drop_useful(self, state, pot_states, player_index):
         - Dish is needed (all pots are full)
         - Nobody is holding a dish
         """
-        if self.num_players != 2: return False
+        if self.num_players != 2:
+            return False
         all_pots_full = len(self.get_full_pots(pot_states)) == self.num_pots
         other_player = state.players[1 - player_index]
-        other_player_holding_dish = other_player.has_object() and other_player.get_object().name == "dish"
+        other_player_holding_dish = (
+            other_player.has_object()
+            and other_player.get_object().name == "dish"
+        )
         return all_pots_full and not other_player_holding_dish
 
     def is_potting_optimal(self, state, old_soup, new_soup):
         """
         True if the highest valued soup possible is the same before and after the potting
         """
-        old_recipe = Recipe(old_soup.ingredients) if old_soup.ingredients else None
+        old_recipe = (
+            Recipe(old_soup.ingredients) if old_soup.ingredients else None
+        )
         new_recipe = Recipe(new_soup.ingredients)
-        old_val = self.get_recipe_value(state, self.get_optimal_possible_recipe(state, old_recipe))
-        new_val = self.get_recipe_value(state, self.get_optimal_possible_recipe(state, new_recipe))
+        old_val = self.get_recipe_value(
+            state, self.get_optimal_possible_recipe(state, old_recipe)
+        )
+        new_val = self.get_recipe_value(
+            state, self.get_optimal_possible_recipe(state, new_recipe)
+        )
         return old_val == new_val
 
     def is_potting_viable(self, state, old_soup, new_soup):
@@ -1761,29 +2257,39 @@ def is_potting_viable(self, state, old_soup, new_soup):
         True if there exists a non-zero reward soup possible from new ingredients
         """
         new_recipe = Recipe(new_soup.ingredients)
-        new_val = self.get_recipe_value(state, self.get_optimal_possible_recipe(state, new_recipe))
+        new_val = self.get_recipe_value(
+            state, self.get_optimal_possible_recipe(state, new_recipe)
+        )
         return new_val > 0
 
     def is_potting_catastrophic(self, state, old_soup, new_soup):
         """
         True if no non-zero reward soup is possible from new ingredients
         """
-        old_recipe = Recipe(old_soup.ingredients) if old_soup.ingredients else None
+        old_recipe = (
+            Recipe(old_soup.ingredients) if old_soup.ingredients else None
+        )
         new_recipe = Recipe(new_soup.ingredients)
-        old_val = self.get_recipe_value(state, self.get_optimal_possible_recipe(state, old_recipe))
-        new_val = self.get_recipe_value(state, self.get_optimal_possible_recipe(state, new_recipe))
+        old_val = self.get_recipe_value(
+            state, self.get_optimal_possible_recipe(state, old_recipe)
+        )
+        new_val = self.get_recipe_value(
+            state, self.get_optimal_possible_recipe(state, new_recipe)
+        )
         return old_val > 0 and new_val == 0
 
     def is_potting_useless(self, state, old_soup, new_soup):
         """
         True if ingredient added to a soup that was already gauranteed to be worth at most 0 points
         """
-        old_recipe = Recipe(old_soup.ingredients) if old_soup.ingredients else None
-        old_val = self.get_recipe_value(state, self.get_optimal_possible_recipe(state, old_recipe))
+        old_recipe = (
+            Recipe(old_soup.ingredients) if old_soup.ingredients else None
+        )
+        old_val = self.get_recipe_value(
+            state, self.get_optimal_possible_recipe(state, old_recipe)
+        )
         return old_val == 0
 
-
-
     #####################
     # TERMINAL GRAPHICS #
     #####################
@@ -1801,7 +2307,11 @@ def state_string(self, state):
                     orientation = player.orientation
                     assert orientation in Direction.ALL_DIRECTIONS
 
-                    player_idx_lst = [i for i, p in enumerate(state.players) if p.position == player.position]
+                    player_idx_lst = [
+                        i
+                        for i, p in enumerate(state.players)
+                        if p.position == player.position
+                    ]
                     assert len(player_idx_lst) == 1
 
                     grid_string_add += Action.ACTION_TO_CHAR[orientation]
@@ -1834,11 +2344,9 @@ def state_string(self, state):
                 grid_string += " "
 
             grid_string += "\n\n"
-        
+
         if state.bonus_orders:
-            grid_string += "Bonus orders: {}\n".format(
-                state.bonus_orders
-            )
+            grid_string += "Bonus orders: {}\n".format(state.bonus_orders)
         # grid_string += "State potential value: {}\n".format(self.potential_function(state))
         return grid_string
 
@@ -1850,40 +2358,72 @@ def state_string(self, state):
     def lossless_state_encoding_shape(self):
         warnings.warn(
             "Using the `lossless_state_encoding_shape` property is deprecated. Please use `get_lossless_state_encoding_shape` method instead",
-            DeprecationWarning
+            DeprecationWarning,
         )
         return np.array(list(self.shape) + [26])
 
     def get_lossless_state_encoding_shape(self):
         return np.array(list(self.shape) + [26])
 
-
-    def lossless_state_encoding(self, overcooked_state, horizon=400, debug=False):
+    def lossless_state_encoding(
+        self, overcooked_state, horizon=400, debug=False
+    ):
         """Featurizes a OvercookedState object into a stack of boolean masks that are easily readable by a CNN"""
-        assert self.num_players == 2, "Functionality has to be added to support encondings for > 2 players"
+        assert (
+            self.num_players == 2
+        ), "Functionality has to be added to support encondings for > 2 players"
         assert type(debug) is bool
-        base_map_features = ["pot_loc", "counter_loc", "onion_disp_loc", "tomato_disp_loc",
-                             "dish_disp_loc", "serve_loc"]
-        variable_map_features = ["onions_in_pot", "tomatoes_in_pot", "onions_in_soup", "tomatoes_in_soup",
-                                 "soup_cook_time_remaining", "soup_done", "dishes", "onions", "tomatoes"]
+        base_map_features = [
+            "pot_loc",
+            "counter_loc",
+            "onion_disp_loc",
+            "tomato_disp_loc",
+            "dish_disp_loc",
+            "serve_loc",
+        ]
+        variable_map_features = [
+            "onions_in_pot",
+            "tomatoes_in_pot",
+            "onions_in_soup",
+            "tomatoes_in_soup",
+            "soup_cook_time_remaining",
+            "soup_done",
+            "dishes",
+            "onions",
+            "tomatoes",
+        ]
         urgency_features = ["urgency"]
         all_objects = overcooked_state.all_objects_list
 
         def make_layer(position, value):
-                layer = np.zeros(self.shape)
-                layer[position] = value
-                return layer
+            layer = np.zeros(self.shape)
+            layer[position] = value
+            return layer
 
         def process_for_player(primary_agent_idx):
             # Ensure that primary_agent_idx layers are ordered before other_agent_idx layers
             other_agent_idx = 1 - primary_agent_idx
-            ordered_player_features = ["player_{}_loc".format(primary_agent_idx), "player_{}_loc".format(other_agent_idx)] + \
-                        ["player_{}_orientation_{}".format(i, Direction.DIRECTION_TO_INDEX[d])
-                        for i, d in itertools.product([primary_agent_idx, other_agent_idx], Direction.ALL_DIRECTIONS)]
+            ordered_player_features = [
+                "player_{}_loc".format(primary_agent_idx),
+                "player_{}_loc".format(other_agent_idx),
+            ] + [
+                "player_{}_orientation_{}".format(
+                    i, Direction.DIRECTION_TO_INDEX[d]
+                )
+                for i, d in itertools.product(
+                    [primary_agent_idx, other_agent_idx],
+                    Direction.ALL_DIRECTIONS,
+                )
+            ]
 
             # LAYERS = ordered_player_features + base_map_features + variable_map_features
-            LAYERS = ordered_player_features + base_map_features + variable_map_features + urgency_features
-            state_mask_dict = {k:np.zeros(self.shape) for k in LAYERS}
+            LAYERS = (
+                ordered_player_features
+                + base_map_features
+                + variable_map_features
+                + urgency_features
+            )
+            state_mask_dict = {k: np.zeros(self.shape) for k in LAYERS}
 
             # MAP LAYERS
             if horizon - overcooked_state.timestep < 40:
@@ -1909,9 +2449,17 @@ def process_for_player(primary_agent_idx):
 
             # PLAYER LAYERS
             for i, player in enumerate(overcooked_state.players):
-                player_orientation_idx = Direction.DIRECTION_TO_INDEX[player.orientation]
-                state_mask_dict["player_{}_loc".format(i)] = make_layer(player.position, 1)
-                state_mask_dict["player_{}_orientation_{}".format(i, player_orientation_idx)] = make_layer(player.position, 1)
+                player_orientation_idx = Direction.DIRECTION_TO_INDEX[
+                    player.orientation
+                ]
+                state_mask_dict["player_{}_loc".format(i)] = make_layer(
+                    player.position, 1
+                )
+                state_mask_dict[
+                    "player_{}_orientation_{}".format(
+                        i, player_orientation_idx
+                    )
+                ] = make_layer(player.position, 1)
 
             # OBJECT & STATE LAYERS
             for obj in all_objects:
@@ -1924,20 +2472,40 @@ def process_for_player(primary_agent_idx):
                     if obj.position in self.get_pot_locations():
                         if obj.is_idle:
                             # onions_in_pot and tomatoes_in_pot are used when the soup is idling, and ingredients could still be added
-                            state_mask_dict["onions_in_pot"] += make_layer(obj.position, ingredients_dict["onion"])
-                            state_mask_dict["tomatoes_in_pot"] += make_layer(obj.position, ingredients_dict["tomato"])
+                            state_mask_dict["onions_in_pot"] += make_layer(
+                                obj.position, ingredients_dict["onion"]
+                            )
+                            state_mask_dict["tomatoes_in_pot"] += make_layer(
+                                obj.position, ingredients_dict["tomato"]
+                            )
                         else:
-                            state_mask_dict["onions_in_soup"] += make_layer(obj.position, ingredients_dict["onion"])
-                            state_mask_dict["tomatoes_in_soup"] += make_layer(obj.position, ingredients_dict["tomato"])
-                            state_mask_dict["soup_cook_time_remaining"] += make_layer(obj.position, obj.cook_time - obj._cooking_tick)
+                            state_mask_dict["onions_in_soup"] += make_layer(
+                                obj.position, ingredients_dict["onion"]
+                            )
+                            state_mask_dict["tomatoes_in_soup"] += make_layer(
+                                obj.position, ingredients_dict["tomato"]
+                            )
+                            state_mask_dict[
+                                "soup_cook_time_remaining"
+                            ] += make_layer(
+                                obj.position, obj.cook_time - obj._cooking_tick
+                            )
                             if obj.is_ready:
-                                state_mask_dict["soup_done"] += make_layer(obj.position, 1)
+                                state_mask_dict["soup_done"] += make_layer(
+                                    obj.position, 1
+                                )
 
                     else:
                         # If player soup is not in a pot, treat it like a soup that is cooked with remaining time 0
-                        state_mask_dict["onions_in_soup"] += make_layer(obj.position, ingredients_dict["onion"])
-                        state_mask_dict["tomatoes_in_soup"] += make_layer(obj.position, ingredients_dict["tomato"])
-                        state_mask_dict["soup_done"] += make_layer(obj.position, 1)
+                        state_mask_dict["onions_in_soup"] += make_layer(
+                            obj.position, ingredients_dict["onion"]
+                        )
+                        state_mask_dict["tomatoes_in_soup"] += make_layer(
+                            obj.position, ingredients_dict["tomato"]
+                        )
+                        state_mask_dict["soup_done"] += make_layer(
+                            obj.position, 1
+                        )
 
                 elif obj.name == "dish":
                     state_mask_dict["dishes"] += make_layer(obj.position, 1)
@@ -1959,7 +2527,9 @@ def process_for_player(primary_agent_idx):
                     print(np.transpose(v, (1, 0)))
 
             # Stack of all the state masks, order decided by order of LAYERS
-            state_mask_stack = np.array([state_mask_dict[layer_id] for layer_id in LAYERS])
+            state_mask_stack = np.array(
+                [state_mask_dict[layer_id] for layer_id in LAYERS]
+            )
             state_mask_stack = np.transpose(state_mask_stack, (1, 2, 0))
             assert state_mask_stack.shape[:2] == self.shape
             assert state_mask_stack.shape[2] == len(LAYERS)
@@ -1968,21 +2538,25 @@ def process_for_player(primary_agent_idx):
 
         # NOTE: Currently not very efficient, a decent amount of computation repeated here
         num_players = len(overcooked_state.players)
-        final_obs_for_players = tuple(process_for_player(i) for i in range(num_players))
+        final_obs_for_players = tuple(
+            process_for_player(i) for i in range(num_players)
+        )
         return final_obs_for_players
 
     @property
     def featurize_state_shape(self):
         warnings.warn(
             "Using the `featurize_state_shape` property is deprecated. Please use `get_featurize_state_shape` method instead",
-            DeprecationWarning
+            DeprecationWarning,
         )
         return self.get_featurize_state_shape(2)
 
     def get_featurize_state_shape(self, num_pots=2):
         num_pot_features = 10
         base_features = 28
-        total_features = self.num_players * (num_pots * num_pot_features + base_features)
+        total_features = self.num_players * (
+            num_pots * num_pot_features + base_features
+        )
         return (total_features,)
 
     def featurize_state(self, overcooked_state, mlam, num_pots=2, **kwargs):
@@ -1992,10 +2566,10 @@ def featurize_state(self, overcooked_state, mlam, num_pots=2, **kwargs):
         Arguments:
             overcooked_state (OvercookedState): state we wish to featurize
             mlam (MediumLevelActionManager): to be used for distance computations necessary for our higher-level feature encodings
-            num_pots (int): Encode the state (ingredients, whether cooking or not, etc) of the 'num_pots' closest pots to each player. 
-                If i < num_pots pots are reachable by player i, then pots [i+1, num_pots] are encoded as all zeros. Changing this 
+            num_pots (int): Encode the state (ingredients, whether cooking or not, etc) of the 'num_pots' closest pots to each player.
+                If i < num_pots pots are reachable by player i, then pots [i+1, num_pots] are encoded as all zeros. Changing this
                 impacts the shape of the feature encoding
-        
+
         Returns:
             ordered_features (list[np.Array]): The ith element contains a player-centric featurized view for the ith player
 
@@ -2018,7 +2592,7 @@ def featurize_state(self, overcooked_state, mlam, num_pots=2, **kwargs):
 
                 other_player_features (length (num_players - 1)*(num_pots*10 + 24)):
                     ordered concatenation of player_{j}_features for j != i
-                
+
                 player_i_dist_to_other_players (length (num_players - 1)*2):
                     [player_j.pos - player_i.pos for j != i]
 
@@ -2042,18 +2616,25 @@ def make_closest_feature(idx, player, name, locations):
                 obj = held_obj
                 feat_dict["p{}_closest_{}".format(i, name)] = (0, 0)
             else:
-                loc, deltas = self.get_deltas_to_closest_location(player, locations, mlam)
+                loc, deltas = self.get_deltas_to_closest_location(
+                    player, locations, mlam
+                )
                 if loc and overcooked_state.has_object(loc):
                     obj = overcooked_state.get_object(loc)
                 feat_dict["p{}_closest_{}".format(idx, name)] = deltas
 
-            if name == 'soup':
+            if name == "soup":
                 num_onions = num_tomatoes = 0
                 if obj:
                     ingredients_cnt = Counter(obj.ingredients)
-                    num_onions, num_tomatoes = ingredients_cnt['onion'], ingredients_cnt['tomato']
+                    num_onions, num_tomatoes = (
+                        ingredients_cnt["onion"],
+                        ingredients_cnt["tomato"],
+                    )
                 feat_dict["p{}_closest_soup_n_onions".format(i)] = [num_onions]
-                feat_dict["p{}_closest_soup_n_tomatoes".format(i)] = [num_tomatoes]
+                feat_dict["p{}_closest_soup_n_tomatoes".format(i)] = [
+                    num_tomatoes
+                ]
 
             return feat_dict
 
@@ -2064,17 +2645,33 @@ def make_pot_feature(idx, player, pot_idx, pot_loc, pot_states):
             # Pot doesn't exist
             feat_dict = {}
             if not pot_loc:
-                feat_dict["p{}_closest_pot_{}_exists".format(idx, pot_idx)] = [0]
-                feat_dict["p{}_closest_pot_{}_is_empty".format(idx, pot_idx)] = [0]
-                feat_dict["p{}_closest_pot_{}_is_full".format(idx, pot_idx)] = [0]
-                feat_dict["p{}_closest_pot_{}_is_cooking".format(idx, pot_idx)] = [0]
-                feat_dict["p{}_closest_pot_{}_is_ready".format(idx, pot_idx)] = [0]
-                feat_dict["p{}_closest_pot_{}_num_onions".format(idx, pot_idx)] = [0]
-                feat_dict["p{}_closest_pot_{}_num_tomatoes".format(idx, pot_idx)] = [0]
-                feat_dict["p{}_closest_pot_{}_cook_time".format(idx, pot_idx)] = [0]
+                feat_dict["p{}_closest_pot_{}_exists".format(idx, pot_idx)] = [
+                    0
+                ]
+                feat_dict[
+                    "p{}_closest_pot_{}_is_empty".format(idx, pot_idx)
+                ] = [0]
+                feat_dict[
+                    "p{}_closest_pot_{}_is_full".format(idx, pot_idx)
+                ] = [0]
+                feat_dict[
+                    "p{}_closest_pot_{}_is_cooking".format(idx, pot_idx)
+                ] = [0]
+                feat_dict[
+                    "p{}_closest_pot_{}_is_ready".format(idx, pot_idx)
+                ] = [0]
+                feat_dict[
+                    "p{}_closest_pot_{}_num_onions".format(idx, pot_idx)
+                ] = [0]
+                feat_dict[
+                    "p{}_closest_pot_{}_num_tomatoes".format(idx, pot_idx)
+                ] = [0]
+                feat_dict[
+                    "p{}_closest_pot_{}_cook_time".format(idx, pot_idx)
+                ] = [0]
                 feat_dict["p{}_closest_pot_{}".format(idx, pot_idx)] = (0, 0)
                 return feat_dict
-            
+
             # Get position information
             deltas = self.get_deltas_to_location(player, pot_loc)
 
@@ -2090,25 +2687,41 @@ def make_pot_feature(idx, player, pot_idx, pot_loc, pot_states):
             if not is_empty:
                 soup = overcooked_state.get_object(pot_loc)
                 ingredients_cnt = Counter(soup.ingredients)
-                num_onions, num_tomatoes = ingredients_cnt['onion'], ingredients_cnt['tomato']
-                cook_time_remaining = 0 if soup.is_idle else soup.cook_time_remaining
-            
+                num_onions, num_tomatoes = (
+                    ingredients_cnt["onion"],
+                    ingredients_cnt["tomato"],
+                )
+                cook_time_remaining = (
+                    0 if soup.is_idle else soup.cook_time_remaining
+                )
+
             # Encode pot and soup info
             feat_dict["p{}_closest_pot_{}_exists".format(idx, pot_idx)] = [1]
-            feat_dict["p{}_closest_pot_{}_is_empty".format(idx, pot_idx)] = [is_empty]
-            feat_dict["p{}_closest_pot_{}_is_full".format(idx, pot_idx)] = [is_full]
-            feat_dict["p{}_closest_pot_{}_is_cooking".format(idx, pot_idx)] = [is_cooking]
-            feat_dict["p{}_closest_pot_{}_is_ready".format(idx, pot_idx)] = [is_ready]
-            feat_dict["p{}_closest_pot_{}_num_onions".format(idx, pot_idx)] = [num_onions]
-            feat_dict["p{}_closest_pot_{}_num_tomatoes".format(idx, pot_idx)] = [num_tomatoes]
-            feat_dict["p{}_closest_pot_{}_cook_time".format(idx, pot_idx)] = [cook_time_remaining]
+            feat_dict["p{}_closest_pot_{}_is_empty".format(idx, pot_idx)] = [
+                is_empty
+            ]
+            feat_dict["p{}_closest_pot_{}_is_full".format(idx, pot_idx)] = [
+                is_full
+            ]
+            feat_dict["p{}_closest_pot_{}_is_cooking".format(idx, pot_idx)] = [
+                is_cooking
+            ]
+            feat_dict["p{}_closest_pot_{}_is_ready".format(idx, pot_idx)] = [
+                is_ready
+            ]
+            feat_dict["p{}_closest_pot_{}_num_onions".format(idx, pot_idx)] = [
+                num_onions
+            ]
+            feat_dict[
+                "p{}_closest_pot_{}_num_tomatoes".format(idx, pot_idx)
+            ] = [num_tomatoes]
+            feat_dict["p{}_closest_pot_{}_cook_time".format(idx, pot_idx)] = [
+                cook_time_remaining
+            ]
             feat_dict["p{}_closest_pot_{}".format(idx, pot_idx)] = deltas
 
             return feat_dict
 
-            
-
-
         IDX_TO_OBJ = ["onion", "soup", "dish", "tomato"]
         OBJ_TO_IDX = {o_name: idx for idx, o_name in enumerate(IDX_TO_OBJ)}
 
@@ -2118,7 +2731,9 @@ def make_pot_feature(idx, player, pot_idx, pot_loc, pot_states):
         for i, player in enumerate(overcooked_state.players):
             # Player info
             orientation_idx = Direction.DIRECTION_TO_INDEX[player.orientation]
-            all_features["p{}_orientation".format(i)] = np.eye(4)[orientation_idx]
+            all_features["p{}_orientation".format(i)] = np.eye(4)[
+                orientation_idx
+            ]
             obj = player.held_object
 
             if obj is None:
@@ -2127,42 +2742,103 @@ def make_pot_feature(idx, player, pot_idx, pot_loc, pot_states):
             else:
                 held_obj_name = obj.name
                 obj_idx = OBJ_TO_IDX[held_obj_name]
-                all_features["p{}_objs".format(i)] = np.eye(len(IDX_TO_OBJ))[obj_idx]
+                all_features["p{}_objs".format(i)] = np.eye(len(IDX_TO_OBJ))[
+                    obj_idx
+                ]
 
             # Closest feature for each object type
-            all_features = concat_dicts(all_features, make_closest_feature(i, player, "onion", self.get_onion_dispenser_locations() + counter_objects["onion"]))
-            all_features = concat_dicts(all_features, make_closest_feature(i, player, "tomato", self.get_tomato_dispenser_locations() + counter_objects["tomato"]))
-            all_features = concat_dicts(all_features, make_closest_feature(i, player, "dish", self.get_dish_dispenser_locations() + counter_objects["dish"]))
-            all_features = concat_dicts(all_features, make_closest_feature(i, player, "soup", counter_objects["soup"]))
-            all_features = concat_dicts(all_features, make_closest_feature(i, player, "serving", self.get_serving_locations()))
-            all_features = concat_dicts(all_features, make_closest_feature(i, player, "empty_counter", self.get_empty_counter_locations(overcooked_state)))
+            all_features = concat_dicts(
+                all_features,
+                make_closest_feature(
+                    i,
+                    player,
+                    "onion",
+                    self.get_onion_dispenser_locations()
+                    + counter_objects["onion"],
+                ),
+            )
+            all_features = concat_dicts(
+                all_features,
+                make_closest_feature(
+                    i,
+                    player,
+                    "tomato",
+                    self.get_tomato_dispenser_locations()
+                    + counter_objects["tomato"],
+                ),
+            )
+            all_features = concat_dicts(
+                all_features,
+                make_closest_feature(
+                    i,
+                    player,
+                    "dish",
+                    self.get_dish_dispenser_locations()
+                    + counter_objects["dish"],
+                ),
+            )
+            all_features = concat_dicts(
+                all_features,
+                make_closest_feature(
+                    i, player, "soup", counter_objects["soup"]
+                ),
+            )
+            all_features = concat_dicts(
+                all_features,
+                make_closest_feature(
+                    i, player, "serving", self.get_serving_locations()
+                ),
+            )
+            all_features = concat_dicts(
+                all_features,
+                make_closest_feature(
+                    i,
+                    player,
+                    "empty_counter",
+                    self.get_empty_counter_locations(overcooked_state),
+                ),
+            )
 
             # Closest pots info
             pot_locations = self.get_pot_locations().copy()
             for pot_idx in range(num_pots):
-                _, closest_pot_loc = mlam.motion_planner.min_cost_to_feature(player.pos_and_or, pot_locations, with_argmin=True)
-                pot_features = make_pot_feature(i, player, pot_idx, closest_pot_loc, pot_states)
+                _, closest_pot_loc = mlam.motion_planner.min_cost_to_feature(
+                    player.pos_and_or, pot_locations, with_argmin=True
+                )
+                pot_features = make_pot_feature(
+                    i, player, pot_idx, closest_pot_loc, pot_states
+                )
                 all_features = concat_dicts(all_features, pot_features)
 
                 if closest_pot_loc:
                     pot_locations.remove(closest_pot_loc)
 
             # Adjacent features info
-            for direction, pos_and_feat in enumerate(self.get_adjacent_features(player)):
+            for direction, pos_and_feat in enumerate(
+                self.get_adjacent_features(player)
+            ):
                 _, feat = pos_and_feat
-                all_features["p{}_wall_{}".format(i, direction)] = [0] if feat == ' ' else [1]
+                all_features["p{}_wall_{}".format(i, direction)] = (
+                    [0] if feat == " " else [1]
+                )
 
         # Convert all list and tuple values to np.arrays
         features_np = {k: np.array(v) for k, v in all_features.items()}
 
-        player_features = [] # Non-position player-specific features
-        player_absolute_positions = [] # Position player-specific features
-        player_relative_positions = [] # Relative position player-specific features
+        player_features = []  # Non-position player-specific features
+        player_absolute_positions = []  # Position player-specific features
+        player_relative_positions = (
+            []
+        )  # Relative position player-specific features
 
         # Compute all player-centric features for each player
         for i, player_i in enumerate(overcooked_state.players):
             # All absolute player-centric features
-            player_i_dict = {k: v for k, v in features_np.items() if k[:2] == "p{}".format(i)}
+            player_i_dict = {
+                k: v
+                for k, v in features_np.items()
+                if k[:2] == "p{}".format(i)
+            }
             features = np.concatenate(list(player_i_dict.values()))
             abs_pos = np.array(player_i.position)
 
@@ -2171,43 +2847,55 @@ def make_pot_feature(idx, player, pot_idx, pot_loc, pot_states):
             for player_j in overcooked_state.players:
                 if player_i == player_j:
                     continue
-                pj_rel_to_pi = np.array(pos_distance(player_j.position, player_i.position))
+                pj_rel_to_pi = np.array(
+                    pos_distance(player_j.position, player_i.position)
+                )
                 rel_pos.append(pj_rel_to_pi)
             rel_pos = np.concatenate(rel_pos)
 
             player_features.append(features)
             player_absolute_positions.append(abs_pos)
             player_relative_positions.append(rel_pos)
-        
+
         # Compute a symmetric, player-centric encoding of features for each player
         ordered_features = []
         for i, player_i in enumerate(overcooked_state.players):
             player_i_features = player_features[i]
             player_i_abs_pos = player_absolute_positions[i]
             player_i_rel_pos = player_relative_positions[i]
-            other_player_features = np.concatenate([feats for j, feats in enumerate(player_features) if j != i])
-            player_i_ordered_features = np.squeeze(np.concatenate([player_i_features, other_player_features, player_i_rel_pos, player_i_abs_pos]))
+            other_player_features = np.concatenate(
+                [feats for j, feats in enumerate(player_features) if j != i]
+            )
+            player_i_ordered_features = np.squeeze(
+                np.concatenate(
+                    [
+                        player_i_features,
+                        other_player_features,
+                        player_i_rel_pos,
+                        player_i_abs_pos,
+                    ]
+                )
+            )
             ordered_features.append(player_i_ordered_features)
 
         return ordered_features
 
-
     def get_deltas_to_closest_location(self, player, locations, mlam):
-        _, closest_loc = mlam.motion_planner.min_cost_to_feature(player.pos_and_or, locations, with_argmin=True)
+        _, closest_loc = mlam.motion_planner.min_cost_to_feature(
+            player.pos_and_or, locations, with_argmin=True
+        )
         deltas = self.get_deltas_to_location(player, closest_loc)
         return closest_loc, deltas
-        
 
     def get_deltas_to_location(self, player, location):
         if location is None:
             # "any object that does not exist or I am carrying is going to show up as a (0,0)
-            # but I can disambiguate the two possibilities by looking at the features 
+            # but I can disambiguate the two possibilities by looking at the features
             # for what kind of object I'm carrying"
             return (0, 0)
         dy_loc, dx_loc = pos_distance(location, player.position)
         return dy_loc, dx_loc
 
-
     ###############################
     # POTENTIAL REWARD SHAPING FN #
     ###############################
@@ -2219,7 +2907,7 @@ def potential_function(self, state, mp, gamma=0.99):
         The main goal here to to approximately infer the actions of an optimal agent, and derive an estimate for the value
         function of the optimal policy. The perfect potential function is indeed the value function
 
-        At a high level, we assume each agent acts independetly, and greedily optimally, and then, using the decay factor "gamma", 
+        At a high level, we assume each agent acts independetly, and greedily optimally, and then, using the decay factor "gamma",
         we calculate the expected discounted reward under this policy
 
         Some implementation details:
@@ -2253,57 +2941,130 @@ def potential_function(self, state, mp, gamma=0.99):
             mp: MotionPlanner instance used to calculate gridworld distances to objects
             gamma: float, discount factor
             max_steps: int, number of steps a high level action is assumed to take in worst case
-        
+
         Returns
             phi(state), the potential of the state
         """
-        if not hasattr(Recipe, '_tomato_value') or not hasattr(Recipe, '_onion_value'):
-            raise ValueError("Potential function requires Recipe onion and tomato values to work properly")
+        if not hasattr(Recipe, "_tomato_value") or not hasattr(
+            Recipe, "_onion_value"
+        ):
+            raise ValueError(
+                "Potential function requires Recipe onion and tomato values to work properly"
+            )
 
         # Constants needed for potential function
         potential_params = {
-            'gamma' : gamma,
-            'tomato_value' : Recipe._tomato_value if Recipe._tomato_value else 13,
-            'onion_value' : Recipe._onion_value if Recipe._tomato_value else 21,
-            **POTENTIAL_CONSTANTS.get(self.layout_name, POTENTIAL_CONSTANTS['default'])
+            "gamma": gamma,
+            "tomato_value": Recipe._tomato_value
+            if Recipe._tomato_value
+            else 13,
+            "onion_value": Recipe._onion_value if Recipe._tomato_value else 21,
+            **POTENTIAL_CONSTANTS.get(
+                self.layout_name, POTENTIAL_CONSTANTS["default"]
+            ),
         }
         pot_states = self.get_pot_states(state)
 
         # Base potential value is the geometric sum of making optimal soups infinitely
-        opt_recipe, discounted_opt_recipe_value = self.get_optimal_possible_recipe(state, None, discounted=True, potential_params=potential_params, return_value=True)
+        (
+            opt_recipe,
+            discounted_opt_recipe_value,
+        ) = self.get_optimal_possible_recipe(
+            state,
+            None,
+            discounted=True,
+            potential_params=potential_params,
+            return_value=True,
+        )
         opt_recipe_value = self.get_recipe_value(state, opt_recipe)
         discount = discounted_opt_recipe_value / opt_recipe_value
         steady_state_value = (discount / (1 - discount)) * opt_recipe_value
         potential = steady_state_value
 
-        # Get list of all soups that have >0 ingredients, sorted based on value of best possible recipe 
-        idle_soups = [state.get_object(pos) for pos in self.get_full_but_not_cooking_pots(pot_states)]
-        idle_soups.extend([state.get_object(pos) for pos in self.get_partially_full_pots(pot_states)])
-        idle_soups = sorted(idle_soups, key=lambda soup : self.get_optimal_possible_recipe(state, Recipe(soup.ingredients), discounted=True, potential_params=potential_params, return_value=True)[1], reverse=True)
+        # Get list of all soups that have >0 ingredients, sorted based on value of best possible recipe
+        idle_soups = [
+            state.get_object(pos)
+            for pos in self.get_full_but_not_cooking_pots(pot_states)
+        ]
+        idle_soups.extend(
+            [
+                state.get_object(pos)
+                for pos in self.get_partially_full_pots(pot_states)
+            ]
+        )
+        idle_soups = sorted(
+            idle_soups,
+            key=lambda soup: self.get_optimal_possible_recipe(
+                state,
+                Recipe(soup.ingredients),
+                discounted=True,
+                potential_params=potential_params,
+                return_value=True,
+            )[1],
+            reverse=True,
+        )
 
         # Build mapping of non_idle soups to the potential value each one will contribue
         # Default potential value is maximimal discount for last two steps applied to optimal recipe value
-        cooking_soups = [state.get_object(pos) for pos in self.get_cooking_pots(pot_states)]
-        done_soups = [state.get_object(pos) for pos in self.get_ready_pots(pot_states)]
-        non_idle_soup_vals = { soup : gamma**(potential_params['max_delivery_steps'] + max(potential_params['max_pickup_steps'], soup.cook_time - soup._cooking_tick)) * max(self.get_recipe_value(state, soup.recipe), 1) for soup in cooking_soups + done_soups }
+        cooking_soups = [
+            state.get_object(pos) for pos in self.get_cooking_pots(pot_states)
+        ]
+        done_soups = [
+            state.get_object(pos) for pos in self.get_ready_pots(pot_states)
+        ]
+        non_idle_soup_vals = {
+            soup: gamma
+            ** (
+                potential_params["max_delivery_steps"]
+                + max(
+                    potential_params["max_pickup_steps"],
+                    soup.cook_time - soup._cooking_tick,
+                )
+            )
+            * max(self.get_recipe_value(state, soup.recipe), 1)
+            for soup in cooking_soups + done_soups
+        }
 
         # Get descriptive list of players based on different attributes
         # Note that these lists are mutually exclusive
-        players_holding_soups = [player for player in state.players if player.has_object() and player.get_object().name == 'soup']
-        players_holding_dishes = [player for player in state.players if player.has_object() and player.get_object().name == 'dish']
-        players_holding_tomatoes = [player for player in state.players if player.has_object() and player.get_object().name == Recipe.TOMATO]
-        players_holding_onions = [player for player in state.players if player.has_object() and player.get_object().name == Recipe.ONION]
-        players_holding_nothing = [player for player in state.players if not player.has_object()]
+        players_holding_soups = [
+            player
+            for player in state.players
+            if player.has_object() and player.get_object().name == "soup"
+        ]
+        players_holding_dishes = [
+            player
+            for player in state.players
+            if player.has_object() and player.get_object().name == "dish"
+        ]
+        players_holding_tomatoes = [
+            player
+            for player in state.players
+            if player.has_object()
+            and player.get_object().name == Recipe.TOMATO
+        ]
+        players_holding_onions = [
+            player
+            for player in state.players
+            if player.has_object() and player.get_object().name == Recipe.ONION
+        ]
+        players_holding_nothing = [
+            player for player in state.players if not player.has_object()
+        ]
 
         ### Step 4 potential ###
 
         # Add potential for each player with a soup
         for player in players_holding_soups:
             # Even if delivery_dist is infinite, we still award potential (as an agent might need to pass the soup to other player first)
-            delivery_dist = mp.min_cost_to_feature(player.pos_and_or, self.terrain_pos_dict['S'])
-            potential += gamma**min(delivery_dist, potential_params['max_delivery_steps']) * max(self.get_recipe_value(state, player.get_object().recipe), 1)
-
-
+            delivery_dist = mp.min_cost_to_feature(
+                player.pos_and_or, self.terrain_pos_dict["S"]
+            )
+            potential += gamma ** min(
+                delivery_dist, potential_params["max_delivery_steps"]
+            ) * max(
+                self.get_recipe_value(state, player.get_object().recipe), 1
+            )
 
         ### Step 3 potential ###
 
@@ -2315,7 +3076,9 @@ def potential_function(self, state, mp, gamma=0.99):
             # find best soup to pick up with dish agent currently has
             for soup in non_idle_soup_vals:
                 # How far away the soup is (inf if not-reachable)
-                pickup_dist = mp.min_cost_to_feature(player.pos_and_or, [soup.position])
+                pickup_dist = mp.min_cost_to_feature(
+                    player.pos_and_or, [soup.position]
+                )
 
                 # mask to award zero score if not reachable
                 # Note: this means that potentially "useful" dish pickups (where agent passes dish to other agent
@@ -2323,9 +3086,14 @@ def potential_function(self, state, mp, gamma=0.99):
                 is_useful = int(pickup_dist < np.inf)
 
                 # Always assume worst-case discounting for step 4, and bump zero-valued soups to 1 as mentioned in docstring
-                pickup_soup_value = gamma**potential_params['max_delivery_steps'] * max(self.get_recipe_value(state, soup.recipe), 1)
+                pickup_soup_value = gamma ** potential_params[
+                    "max_delivery_steps"
+                ] * max(self.get_recipe_value(state, soup.recipe), 1)
                 cook_time_remaining = soup.cook_time - soup._cooking_tick
-                discount = gamma**max(cook_time_remaining, min(pickup_dist, potential_params['max_pickup_steps']))
+                discount = gamma ** max(
+                    cook_time_remaining,
+                    min(pickup_dist, potential_params["max_pickup_steps"]),
+                )
 
                 # Final discount-adjusted value for this player pursuing this soup
                 pickup_value = discount * pickup_soup_value * is_useful
@@ -2334,25 +3102,30 @@ def potential_function(self, state, mp, gamma=0.99):
                 if pickup_dist < np.inf and pickup_value > best_pickup_value:
                     best_pickup_soup = soup
                     best_pickup_value = pickup_value
-            
+
             # Set best-case score for this soup. Can only improve upon previous players policies
             # Note cooperative policies between players not considered
             if best_pickup_soup:
-                non_idle_soup_vals[best_pickup_soup] = max(non_idle_soup_vals[best_pickup_soup], best_pickup_value)
+                non_idle_soup_vals[best_pickup_soup] = max(
+                    non_idle_soup_vals[best_pickup_soup], best_pickup_value
+                )
 
         # Apply potential for each idle soup as calculated above
         for soup in non_idle_soup_vals:
             potential += non_idle_soup_vals[soup]
 
-
-
         ### Step 2 potential ###
 
         # Iterate over idle soups in decreasing order of value so we greedily prioritize higher valued soups
         for soup in idle_soups:
             # Calculate optimal recipe
             curr_recipe = Recipe(soup.ingredients)
-            opt_recipe = self.get_optimal_possible_recipe(state, curr_recipe, discounted=True, potential_params=potential_params)
+            opt_recipe = self.get_optimal_possible_recipe(
+                state,
+                curr_recipe,
+                discounted=True,
+                potential_params=potential_params,
+            )
 
             # Calculate missing ingredients needed to complete optimal recipe
             missing_ingredients = list(opt_recipe.ingredients)
@@ -2360,25 +3133,36 @@ def potential_function(self, state, mp, gamma=0.99):
                 missing_ingredients.remove(ingredient)
 
             # Base discount for steps 3-4
-            discount = gamma**(max(potential_params['max_pickup_steps'], opt_recipe.time) + potential_params['max_delivery_steps'])
+            discount = gamma ** (
+                max(potential_params["max_pickup_steps"], opt_recipe.time)
+                + potential_params["max_delivery_steps"]
+            )
 
             # Add a multiplicative discount for each needed ingredient (this has the effect of giving more award to soups
             # that are closer to being completed)
             for ingredient in missing_ingredients:
                 # Players who might have an ingredient we need
-                pertinent_players = players_holding_tomatoes if ingredient == Recipe.TOMATO else players_holding_onions
+                pertinent_players = (
+                    players_holding_tomatoes
+                    if ingredient == Recipe.TOMATO
+                    else players_holding_onions
+                )
                 dist = np.inf
                 closest_player = None
 
                 # Find closest player with ingredient we need
                 for player in pertinent_players:
-                    curr_dist = mp.min_cost_to_feature(player.pos_and_or, [soup.position])
+                    curr_dist = mp.min_cost_to_feature(
+                        player.pos_and_or, [soup.position]
+                    )
                     if curr_dist < dist:
                         dist = curr_dist
                         closest_player = player
 
                 # Update discount to account for adding this missing ingredient (defaults to min_coeff if no pertinent players exist)
-                discount *= gamma**min(dist, potential_params['pot_{}_steps'.format(ingredient)])
+                discount *= gamma ** min(
+                    dist, potential_params["pot_{}_steps".format(ingredient)]
+                )
 
                 # Cross off this player's ingreident contribution so it can't be double-counted
                 if closest_player:
@@ -2391,28 +3175,59 @@ def potential_function(self, state, mp, gamma=0.99):
                 discount *= gamma
             else:
                 # Otherwise, we assume that every player holding nothing will make a beeline to this soup since it's already optimal
-                cook_dist = min([mp.min_cost_to_feature(player.pos_and_or, [soup.position]) for player in players_holding_nothing], default=np.inf)
-                discount *= gamma**min(cook_dist, potential_params['max_pickup_steps'])
-
-            potential += discount * max(self.get_recipe_value(state, opt_recipe), 1)
-
+                cook_dist = min(
+                    [
+                        mp.min_cost_to_feature(
+                            player.pos_and_or, [soup.position]
+                        )
+                        for player in players_holding_nothing
+                    ],
+                    default=np.inf,
+                )
+                discount *= gamma ** min(
+                    cook_dist, potential_params["max_pickup_steps"]
+                )
+
+            potential += discount * max(
+                self.get_recipe_value(state, opt_recipe), 1
+            )
 
         ### Step 1 Potential ###
 
         # Add potential for each tomato that is left over after using all others to complete optimal recipes
         for player in players_holding_tomatoes:
             # will be inf if there exists no empty pot that is reachable
-            dist = mp.min_cost_to_feature(player.pos_and_or, self.get_empty_pots(pot_states))
+            dist = mp.min_cost_to_feature(
+                player.pos_and_or, self.get_empty_pots(pot_states)
+            )
             is_useful = int(dist < np.inf)
-            discount = gamma**(min(potential_params['pot_tomato_steps'], dist) + potential_params['max_pickup_steps'] + potential_params['max_delivery_steps']) * is_useful
-            potential += discount * potential_params['tomato_value']
+            discount = (
+                gamma
+                ** (
+                    min(potential_params["pot_tomato_steps"], dist)
+                    + potential_params["max_pickup_steps"]
+                    + potential_params["max_delivery_steps"]
+                )
+                * is_useful
+            )
+            potential += discount * potential_params["tomato_value"]
 
         # Add potential for each onion that is remaining after using others to complete optimal recipes if possible
         for player in players_holding_onions:
-            dist = mp.min_cost_to_feature(player.pos_and_or, self.get_empty_pots(pot_states))
+            dist = mp.min_cost_to_feature(
+                player.pos_and_or, self.get_empty_pots(pot_states)
+            )
             is_useful = int(dist < np.inf)
-            discount = gamma**(min(potential_params['pot_onion_steps'], dist) + potential_params['max_pickup_steps'] + potential_params['max_delivery_steps']) * is_useful
-            potential += discount * potential_params['onion_value']
+            discount = (
+                gamma
+                ** (
+                    min(potential_params["pot_onion_steps"], dist)
+                    + potential_params["max_pickup_steps"]
+                    + potential_params["max_delivery_steps"]
+                )
+                * is_useful
+            )
+            potential += discount * potential_params["onion_value"]
 
         # At last
         return potential
diff --git a/src/overcooked_ai_py/mdp/overcooked_trajectory.py b/src/overcooked_ai_py/mdp/overcooked_trajectory.py
index 158a3b24..3aa80416 100644
--- a/src/overcooked_ai_py/mdp/overcooked_trajectory.py
+++ b/src/overcooked_ai_py/mdp/overcooked_trajectory.py
@@ -31,12 +31,20 @@
     }
 """
 
-TIMESTEP_TRAJ_KEYS = set(["ep_states", "ep_actions", "ep_rewards", "ep_dones", "ep_infos"])
-EPISODE_TRAJ_KEYS = set(["ep_returns", "ep_lengths", "mdp_params", "env_params"])
-DEFAULT_TRAJ_KEYS = set(list(TIMESTEP_TRAJ_KEYS) + list(EPISODE_TRAJ_KEYS) + ["metadatas"])
+TIMESTEP_TRAJ_KEYS = set(
+    ["ep_states", "ep_actions", "ep_rewards", "ep_dones", "ep_infos"]
+)
+EPISODE_TRAJ_KEYS = set(
+    ["ep_returns", "ep_lengths", "mdp_params", "env_params"]
+)
+DEFAULT_TRAJ_KEYS = set(
+    list(TIMESTEP_TRAJ_KEYS) + list(EPISODE_TRAJ_KEYS) + ["metadatas"]
+)
+
 
 def get_empty_trajectory():
-    return { k : [] if k != 'metadatas' else {} for k in DEFAULT_TRAJ_KEYS }
+    return {k: [] if k != "metadatas" else {} for k in DEFAULT_TRAJ_KEYS}
+
 
 def append_trajectories(traj_one, traj_two):
     # Note: Drops metadatas for now
@@ -46,20 +54,27 @@ def append_trajectories(traj_one, traj_two):
         traj_one = get_empty_trajectory()
     if not traj_two:
         traj_two = get_empty_trajectory()
-    
-    if set(traj_one.keys()) != DEFAULT_TRAJ_KEYS or set(traj_two.keys()) != DEFAULT_TRAJ_KEYS:
+
+    if (
+        set(traj_one.keys()) != DEFAULT_TRAJ_KEYS
+        or set(traj_two.keys()) != DEFAULT_TRAJ_KEYS
+    ):
         raise ValueError("Trajectory key mismatch!")
-    
-    appended_traj = { "metadatas" : {} }
+
+    appended_traj = {"metadatas": {}}
     for k in traj_one:
-        if k != 'metadatas':
+        if k != "metadatas":
             traj_one_value = traj_one[k]
             traj_two_value = traj_two[k]
-            assert type(traj_one_value) == type(traj_two_value), "mismatched trajectory types!"
-            
+            assert type(traj_one_value) == type(
+                traj_two_value
+            ), "mismatched trajectory types!"
+
             if type(traj_one_value) == list:
                 appended_traj[k] = traj_one_value + traj_two_value
             else:
-                appended_traj[k] = np.concatenate([traj_one_value, traj_two_value], axis=0)
+                appended_traj[k] = np.concatenate(
+                    [traj_one_value, traj_two_value], axis=0
+                )
 
-    return appended_traj
\ No newline at end of file
+    return appended_traj
diff --git a/src/overcooked_ai_py/planning/planners.py b/src/overcooked_ai_py/planning/planners.py
index d226b707..d701848f 100644
--- a/src/overcooked_ai_py/planning/planners.py
+++ b/src/overcooked_ai_py/planning/planners.py
@@ -1,37 +1,50 @@
-import itertools, os
+import itertools
+import os
+import pickle
+import time
+
 import numpy as np
-import pickle, time
-from overcooked_ai_py.utils import manhattan_distance
-from overcooked_ai_py.planning.search import Graph, NotConnectedError
+
+from overcooked_ai_py.data.planners import (
+    PLANNERS_DIR,
+    load_saved_action_manager,
+    load_saved_motion_planner,
+)
 from overcooked_ai_py.mdp.actions import Action, Direction
-from overcooked_ai_py.mdp.overcooked_mdp import OvercookedState, PlayerState, OvercookedGridworld, EVENT_TYPES
-from overcooked_ai_py.data.planners import load_saved_action_manager, load_saved_motion_planner, PLANNERS_DIR
+from overcooked_ai_py.mdp.overcooked_mdp import (
+    EVENT_TYPES,
+    OvercookedGridworld,
+    OvercookedState,
+    PlayerState,
+)
+from overcooked_ai_py.planning.search import Graph, NotConnectedError
+from overcooked_ai_py.utils import manhattan_distance
 
 # Run planning logic with additional checks and
 # computation to prevent or identify possible minor errors
 SAFE_RUN = False
 
 NO_COUNTERS_PARAMS = {
-        'start_orientations': False,
-        'wait_allowed': False,
-        'counter_goals': [],
-        'counter_drop': [],
-        'counter_pickup': [],
-        'same_motion_goals': True
+    "start_orientations": False,
+    "wait_allowed": False,
+    "counter_goals": [],
+    "counter_drop": [],
+    "counter_pickup": [],
+    "same_motion_goals": True,
 }
 
 NO_COUNTERS_START_OR_PARAMS = {
-    'start_orientations': True,
-    'wait_allowed': False,
-    'counter_goals': [],
-    'counter_drop': [],
-    'counter_pickup': [],
-    'same_motion_goals': True
+    "start_orientations": True,
+    "wait_allowed": False,
+    "counter_goals": [],
+    "counter_drop": [],
+    "counter_pickup": [],
+    "same_motion_goals": True,
 }
 
 
 class MotionPlanner(object):
-    """A planner that computes optimal plans for a single agent to 
+    """A planner that computes optimal plans for a single agent to
     arrive at goal positions and orientations in an OvercookedGridworld.
 
     Args:
@@ -43,7 +56,7 @@ class MotionPlanner(object):
     def __init__(self, mdp, counter_goals=[]):
         self.mdp = mdp
 
-        # If positions facing counters should be 
+        # If positions facing counters should be
         # allowed as motion goals
         self.counter_goals = counter_goals
 
@@ -55,7 +68,7 @@ def __init__(self, mdp, counter_goals=[]):
         self.all_plans = self._populate_all_plans()
 
     def save_to_file(self, filename):
-        with open(filename, 'wb') as output:
+        with open(filename, "wb") as output:
             pickle.dump(self, output, pickle.HIGHEST_PROTOCOL)
 
     @staticmethod
@@ -63,10 +76,20 @@ def from_file(filename):
         return load_saved_motion_planner(filename)
 
     @staticmethod
-    def from_pickle_or_compute(mdp, counter_goals, custom_filename=None, force_compute=False, info=False):
+    def from_pickle_or_compute(
+        mdp,
+        counter_goals,
+        custom_filename=None,
+        force_compute=False,
+        info=False,
+    ):
         assert isinstance(mdp, OvercookedGridworld)
 
-        filename = custom_filename if custom_filename is not None else mdp.layout_name + "_mp.pkl"
+        filename = (
+            custom_filename
+            if custom_filename is not None
+            else mdp.layout_name + "_mp.pkl"
+        )
 
         if force_compute:
             return MotionPlanner.compute_mp(filename, mdp, counter_goals)
@@ -76,34 +99,48 @@ def from_pickle_or_compute(mdp, counter_goals, custom_filename=None, force_compu
 
             if mp.counter_goals != counter_goals or mp.mdp != mdp:
                 if info:
-                    print("motion planner with different counter goal or mdp found, computing from scratch")
+                    print(
+                        "motion planner with different counter goal or mdp found, computing from scratch"
+                    )
                 return MotionPlanner.compute_mp(filename, mdp, counter_goals)
 
-        except (FileNotFoundError, ModuleNotFoundError, EOFError, AttributeError) as e:
+        except (
+            FileNotFoundError,
+            ModuleNotFoundError,
+            EOFError,
+            AttributeError,
+        ) as e:
             if info:
                 print("Recomputing motion planner due to:", e)
             return MotionPlanner.compute_mp(filename, mdp, counter_goals)
 
         if info:
-            print("Loaded MotionPlanner from {}".format(os.path.join(PLANNERS_DIR, filename)))
+            print(
+                "Loaded MotionPlanner from {}".format(
+                    os.path.join(PLANNERS_DIR, filename)
+                )
+            )
         return mp
 
     @staticmethod
     def compute_mp(filename, mdp, counter_goals):
         final_filepath = os.path.join(PLANNERS_DIR, filename)
-        print("Computing MotionPlanner to be saved in {}".format(final_filepath))
+        print(
+            "Computing MotionPlanner to be saved in {}".format(final_filepath)
+        )
         start_time = time.time()
         mp = MotionPlanner(mdp, counter_goals)
-        print("It took {} seconds to create mp".format(time.time() - start_time))
+        print(
+            "It took {} seconds to create mp".format(time.time() - start_time)
+        )
         mp.save_to_file(final_filepath)
         return mp
 
-
     def get_plan(self, start_pos_and_or, goal_pos_and_or):
         """
         Returns pre-computed plan from initial agent position
         and orientation to a goal position and orientation.
-        
+
         Args:
             start_pos_and_or (tuple): starting (pos, or) tuple
             goal_pos_and_or (tuple): goal (pos, or) tuple
@@ -111,20 +148,21 @@ def get_plan(self, start_pos_and_or, goal_pos_and_or):
         plan_key = (start_pos_and_or, goal_pos_and_or)
         action_plan, pos_and_or_path, plan_cost = self.all_plans[plan_key]
         return action_plan, pos_and_or_path, plan_cost
-    
+
     def get_gridworld_distance(self, start_pos_and_or, goal_pos_and_or):
         """Number of actions necessary to go from starting position
         and orientations to goal position and orientation (not including
         interaction action)"""
-        assert self.is_valid_motion_start_goal_pair(start_pos_and_or, goal_pos_and_or), \
-            "Goal position and orientation were not a valid motion goal"
+        assert self.is_valid_motion_start_goal_pair(
+            start_pos_and_or, goal_pos_and_or
+        ), "Goal position and orientation were not a valid motion goal"
         _, _, plan_cost = self.get_plan(start_pos_and_or, goal_pos_and_or)
         # Removing interaction cost
         return plan_cost - 1
 
     def get_gridworld_pos_distance(self, pos1, pos2):
-        """Minimum (over possible orientations) number of actions necessary 
-        to go from starting position to goal position (not including 
+        """Minimum (over possible orientations) number of actions necessary
+        to go from starting position to goal position (not including
         interaction action)."""
         # NOTE: currently unused, pretty bad code. If used in future, clean up
         min_cost = np.Inf
@@ -140,17 +178,29 @@ def get_gridworld_pos_distance(self, pos1, pos2):
     def _populate_all_plans(self):
         """Pre-computes all valid plans from any valid pos_or to any valid motion_goal"""
         all_plans = {}
-        valid_pos_and_ors = self.mdp.get_valid_player_positions_and_orientations()
-        valid_motion_goals = filter(self.is_valid_motion_goal, valid_pos_and_ors)
-        for start_motion_state, goal_motion_state in itertools.product(valid_pos_and_ors, valid_motion_goals):
-            if not self.is_valid_motion_start_goal_pair(start_motion_state, goal_motion_state):
+        valid_pos_and_ors = (
+            self.mdp.get_valid_player_positions_and_orientations()
+        )
+        valid_motion_goals = filter(
+            self.is_valid_motion_goal, valid_pos_and_ors
+        )
+        for start_motion_state, goal_motion_state in itertools.product(
+            valid_pos_and_ors, valid_motion_goals
+        ):
+            if not self.is_valid_motion_start_goal_pair(
+                start_motion_state, goal_motion_state
+            ):
                 continue
-            action_plan, pos_and_or_path, plan_cost = self._compute_plan(start_motion_state, goal_motion_state)
+            action_plan, pos_and_or_path, plan_cost = self._compute_plan(
+                start_motion_state, goal_motion_state
+            )
             plan_key = (start_motion_state, goal_motion_state)
             all_plans[plan_key] = (action_plan, pos_and_or_path, plan_cost)
         return all_plans
 
-    def is_valid_motion_start_goal_pair(self, start_pos_and_or, goal_pos_and_or):
+    def is_valid_motion_start_goal_pair(
+        self, start_pos_and_or, goal_pos_and_or
+    ):
         if not self.is_valid_motion_goal(goal_pos_and_or):
             return False
         # the valid motion start goal needs to be in the same connected component
@@ -159,33 +209,52 @@ def is_valid_motion_start_goal_pair(self, start_pos_and_or, goal_pos_and_or):
         return True
 
     def is_valid_motion_goal(self, goal_pos_and_or):
-        """Checks that desired single-agent goal state (position and orientation) 
+        """Checks that desired single-agent goal state (position and orientation)
         is reachable and is facing a terrain feature"""
         goal_position, goal_orientation = goal_pos_and_or
         if goal_position not in self.mdp.get_valid_player_positions():
             return False
 
         # Restricting goals to be facing a terrain feature
-        pos_of_facing_terrain = Action.move_in_direction(goal_position, goal_orientation)
-        facing_terrain_type = self.mdp.get_terrain_type_at_pos(pos_of_facing_terrain)
-        if facing_terrain_type == ' ' or (facing_terrain_type == 'X' and pos_of_facing_terrain not in self.counter_goals):
+        pos_of_facing_terrain = Action.move_in_direction(
+            goal_position, goal_orientation
+        )
+        facing_terrain_type = self.mdp.get_terrain_type_at_pos(
+            pos_of_facing_terrain
+        )
+        if facing_terrain_type == " " or (
+            facing_terrain_type == "X"
+            and pos_of_facing_terrain not in self.counter_goals
+        ):
             return False
         return True
 
     def _compute_plan(self, start_motion_state, goal_motion_state):
         """Computes optimal action plan for single agent movement
-        
+
         Args:
             start_motion_state (tuple): starting positions and orientations
             goal_motion_state (tuple): goal positions and orientations
         """
-        assert self.is_valid_motion_start_goal_pair(start_motion_state, goal_motion_state)
-        positions_plan = self._get_position_plan_from_graph(start_motion_state, goal_motion_state)
-        action_plan, pos_and_or_path, plan_length = self.action_plan_from_positions(positions_plan, start_motion_state, goal_motion_state)
+        assert self.is_valid_motion_start_goal_pair(
+            start_motion_state, goal_motion_state
+        )
+        positions_plan = self._get_position_plan_from_graph(
+            start_motion_state, goal_motion_state
+        )
+        (
+            action_plan,
+            pos_and_or_path,
+            plan_length,
+        ) = self.action_plan_from_positions(
+            positions_plan, start_motion_state, goal_motion_state
+        )
         return action_plan, pos_and_or_path, plan_length
 
     def positions_are_connected(self, start_pos_and_or, goal_pos_and_or):
-        return self.graph_problem.are_in_same_cc(start_pos_and_or, goal_pos_and_or)
+        return self.graph_problem.are_in_same_cc(
+            start_pos_and_or, goal_pos_and_or
+        )
 
     def _get_position_plan_from_graph(self, start_node, end_node):
         """Recovers positions to be reached by agent after the start node to reach the end node"""
@@ -194,11 +263,13 @@ def _get_position_plan_from_graph(self, start_node, end_node):
         positions_plan = [state_node[0] for state_node in node_path[1:]]
         return positions_plan
 
-    def action_plan_from_positions(self, position_list, start_motion_state, goal_motion_state):
+    def action_plan_from_positions(
+        self, position_list, start_motion_state, goal_motion_state
+    ):
         """
         Recovers an action plan reaches the goal motion position and orientation, and executes
         and interact action.
-        
+
         Args:
             position_list (list): list of positions to be reached after the starting position
                                   (does not include starting position, but includes ending position)
@@ -218,15 +289,19 @@ def action_plan_from_positions(self, position_list, start_motion_state, goal_mot
         # Get agent to goal position
         while position_to_go and curr_pos != goal_position:
             next_pos = position_to_go.pop(0)
-            action = Action.determine_action_for_change_in_pos(curr_pos, next_pos)
+            action = Action.determine_action_for_change_in_pos(
+                curr_pos, next_pos
+            )
             action_plan.append(action)
             curr_or = action if action != Action.STAY else curr_or
             pos_and_or_path.append((next_pos, curr_or))
             curr_pos = next_pos
-        
+
         # Fix agent orientation if necessary
         if curr_or != goal_orientation:
-            new_pos, _ = self.mdp._move_if_direction(curr_pos, curr_or, goal_orientation)
+            new_pos, _ = self.mdp._move_if_direction(
+                curr_pos, curr_or, goal_orientation
+            )
             assert new_pos == goal_position
             action_plan.append(goal_orientation)
             pos_and_or_path.append((goal_position, goal_orientation))
@@ -240,17 +315,27 @@ def action_plan_from_positions(self, position_list, start_motion_state, goal_mot
     def _graph_from_grid(self):
         """Creates a graph adjacency matrix from an Overcooked MDP class."""
         state_decoder = {}
-        for state_index, motion_state in enumerate(self.mdp.get_valid_player_positions_and_orientations()):
+        for state_index, motion_state in enumerate(
+            self.mdp.get_valid_player_positions_and_orientations()
+        ):
             state_decoder[state_index] = motion_state
 
-        pos_encoder = {motion_state:state_index for state_index, motion_state in state_decoder.items()}
+        pos_encoder = {
+            motion_state: state_index
+            for state_index, motion_state in state_decoder.items()
+        }
         num_graph_nodes = len(state_decoder)
 
         adjacency_matrix = np.zeros((num_graph_nodes, num_graph_nodes))
         for state_index, start_motion_state in state_decoder.items():
-            for action, successor_motion_state in self._get_valid_successor_motion_states(start_motion_state):
+            for (
+                action,
+                successor_motion_state,
+            ) in self._get_valid_successor_motion_states(start_motion_state):
                 adj_pos_index = pos_encoder[successor_motion_state]
-                adjacency_matrix[state_index][adj_pos_index] = self._graph_action_cost(action)
+                adjacency_matrix[state_index][
+                    adj_pos_index
+                ] = self._graph_action_cost(action)
 
         return Graph(adjacency_matrix, pos_encoder, state_decoder)
 
@@ -262,17 +347,30 @@ def _graph_action_cost(self, action):
     def _get_valid_successor_motion_states(self, start_motion_state):
         """Get valid motion states one action away from the starting motion state."""
         start_position, start_orientation = start_motion_state
-        return [(action, self.mdp._move_if_direction(start_position, start_orientation, action)) for action in Action.ALL_ACTIONS]
+        return [
+            (
+                action,
+                self.mdp._move_if_direction(
+                    start_position, start_orientation, action
+                ),
+            )
+            for action in Action.ALL_ACTIONS
+        ]
 
-    def min_cost_between_features(self, pos_list1, pos_list2, manhattan_if_fail=False):
+    def min_cost_between_features(
+        self, pos_list1, pos_list2, manhattan_if_fail=False
+    ):
         """
-        Determines the minimum number of timesteps necessary for a player to go from any 
+        Determines the minimum number of timesteps necessary for a player to go from any
         terrain feature in list1 to any feature in list2 and perform an interact action
         """
         min_dist = np.Inf
         min_manhattan = np.Inf
         for pos1, pos2 in itertools.product(pos_list1, pos_list2):
-            for mg1, mg2 in itertools.product(self.motion_goals_for_pos[pos1], self.motion_goals_for_pos[pos2]):
+            for mg1, mg2 in itertools.product(
+                self.motion_goals_for_pos[pos1],
+                self.motion_goals_for_pos[pos2],
+            ):
                 if not self.is_valid_motion_start_goal_pair(mg1, mg2):
                     if manhattan_if_fail:
                         pos0, pos1 = mg1[0], mg2[0]
@@ -283,27 +381,37 @@ def min_cost_between_features(self, pos_list1, pos_list2, manhattan_if_fail=Fals
                 curr_dist = self.get_gridworld_distance(mg1, mg2)
                 if curr_dist < min_dist:
                     min_dist = curr_dist
-                
+
         # +1 to account for interaction action
         if manhattan_if_fail and min_dist == np.Inf:
             min_dist = min_manhattan
         min_cost = min_dist + 1
         return min_cost
 
-    def min_cost_to_feature(self, start_pos_and_or, feature_pos_list, with_argmin=False, debug=False):
+    def min_cost_to_feature(
+        self,
+        start_pos_and_or,
+        feature_pos_list,
+        with_argmin=False,
+        debug=False,
+    ):
         """
         Determines the minimum number of timesteps necessary for a player to go from the starting
         position and orientation to any feature in feature_pos_list and perform an interact action
         """
         start_pos = start_pos_and_or[0]
-        assert self.mdp.get_terrain_type_at_pos(start_pos) != 'X'
+        assert self.mdp.get_terrain_type_at_pos(start_pos) != "X"
         min_dist = np.Inf
         best_feature = None
         for feature_pos in feature_pos_list:
             for feature_goal in self.motion_goals_for_pos[feature_pos]:
-                if not self.is_valid_motion_start_goal_pair(start_pos_and_or, feature_goal):
+                if not self.is_valid_motion_start_goal_pair(
+                    start_pos_and_or, feature_goal
+                ):
                     continue
-                curr_dist = self.get_gridworld_distance(start_pos_and_or, feature_goal)
+                curr_dist = self.get_gridworld_distance(
+                    start_pos_and_or, feature_goal
+                )
                 if curr_dist < min_dist:
                     best_feature = feature_pos
                     min_dist = curr_dist
@@ -319,9 +427,14 @@ def _get_goal_dict(self):
         terrain features that the agent might want to interact with."""
         terrain_feature_locations = []
         for terrain_type, pos_list in self.mdp.terrain_pos_dict.items():
-            if terrain_type != ' ':
+            if terrain_type != " ":
                 terrain_feature_locations += pos_list
-        return {feature_pos:self._get_possible_motion_goals_for_feature(feature_pos) for feature_pos in terrain_feature_locations}
+        return {
+            feature_pos: self._get_possible_motion_goals_for_feature(
+                feature_pos
+            )
+            for feature_pos in terrain_feature_locations
+        }
 
     def _get_possible_motion_goals_for_feature(self, goal_pos):
         """Returns a list of possible goal positions (and orientations)
@@ -338,7 +451,7 @@ def _get_possible_motion_goals_for_feature(self, goal_pos):
 
 
 class JointMotionPlanner(object):
-    """A planner that computes optimal plans for a two agents to 
+    """A planner that computes optimal plans for a two agents to
     arrive at goal positions and orientations in a OvercookedGridworld.
 
     Args:
@@ -349,7 +462,7 @@ def __init__(self, mdp, params, debug=False):
         self.mdp = mdp
 
         # Whether starting orientations should be accounted for
-        # when solving all motion problems 
+        # when solving all motion problems
         # (increases number of plans by a factor of 4)
         # but removes additional fudge factor <= 1 for each
         # joint motion plan
@@ -358,11 +471,13 @@ def __init__(self, mdp, params, debug=False):
 
         # Enable both agents to have the same motion goal
         self.same_motion_goals = params["same_motion_goals"]
-        
+
         # Single agent motion planner
-        self.motion_planner = MotionPlanner(mdp, counter_goals=params["counter_goals"])
+        self.motion_planner = MotionPlanner(
+            mdp, counter_goals=params["counter_goals"]
+        )
 
-        # Graph problem that returns optimal paths from 
+        # Graph problem that returns optimal paths from
         # starting positions to goal positions (without
         # accounting for orientations)
         self.joint_graph_problem = self._joint_graph_from_grid()
@@ -372,30 +487,37 @@ def get_low_level_action_plan(self, start_jm_state, goal_jm_state):
         """
         Returns pre-computed plan from initial joint motion state
         to a goal joint motion state.
-        
+
         Args:
             start_jm_state (tuple): starting pos & orients ((pos1, or1), (pos2, or2))
             goal_jm_state (tuple): goal pos & orients ((pos1, or1), (pos2, or2))
-        
+
         Returns:
             joint_action_plan (list): joint actions to be executed to reach end_jm_state
             end_jm_state (tuple): the pair of (pos, or) tuples corresponding
-                to the ending timestep (this will usually be different from 
+                to the ending timestep (this will usually be different from
                 goal_jm_state, as one agent will end before other).
             plan_lengths (tuple): lengths for each agent's plan
         """
-        assert self.is_valid_joint_motion_pair(start_jm_state, goal_jm_state), \
-            "start: {} \t end: {} was not a valid motion goal pair".format(start_jm_state, goal_jm_state)
+        assert self.is_valid_joint_motion_pair(
+            start_jm_state, goal_jm_state
+        ), "start: {} \t end: {} was not a valid motion goal pair".format(
+            start_jm_state, goal_jm_state
+        )
 
         if self.start_orientations:
             plan_key = (start_jm_state, goal_jm_state)
         else:
-            starting_positions = tuple(player_pos_and_or[0] for player_pos_and_or in start_jm_state)
-            goal_positions = tuple(player_pos_and_or[0] for player_pos_and_or in goal_jm_state)
+            starting_positions = tuple(
+                player_pos_and_or[0] for player_pos_and_or in start_jm_state
+            )
+            goal_positions = tuple(
+                player_pos_and_or[0] for player_pos_and_or in goal_jm_state
+            )
             # If beginning positions are equal to end positions, the pre-stored
-            # plan (not dependent on initial positions) will likely return a 
+            # plan (not dependent on initial positions) will likely return a
             # wrong answer, so we compute it from scratch.
-            # 
+            #
             # This is because we only compute plans with starting orientations
             # (North, North), so if one of the two agents starts at location X
             # with orientation East it's goal is to get to location X with
@@ -406,17 +528,23 @@ def get_low_level_action_plan(self, start_jm_state, goal_jm_state):
             # We also compute the plan for any shared motion goal with SAFE_RUN,
             # as there are some minor edge cases that could not be accounted for
             # but I expect should not make a difference in nearly all scenarios
-            if any([s == g for s, g in zip(starting_positions, goal_positions)]) or (SAFE_RUN and goal_positions[0] == goal_positions[1]):
+            if any(
+                [s == g for s, g in zip(starting_positions, goal_positions)]
+            ) or (SAFE_RUN and goal_positions[0] == goal_positions[1]):
                 return self._obtain_plan(start_jm_state, goal_jm_state)
 
             dummy_orientation = Direction.NORTH
-            dummy_start_jm_state = tuple((pos, dummy_orientation) for pos in starting_positions)
+            dummy_start_jm_state = tuple(
+                (pos, dummy_orientation) for pos in starting_positions
+            )
             plan_key = (dummy_start_jm_state, goal_jm_state)
 
         if plan_key not in self.all_plans:
             num_player = len(goal_jm_state)
             return [], None, [np.inf] * num_player
-        joint_action_plan, end_jm_state, plan_lengths = self.all_plans[plan_key]
+        joint_action_plan, end_jm_state, plan_lengths = self.all_plans[
+            plan_key
+        ]
         return joint_action_plan, end_jm_state, plan_lengths
 
     def _populate_all_plans(self):
@@ -425,31 +553,57 @@ def _populate_all_plans(self):
 
         # Joint states are valid if players are not in same location
         if self.start_orientations:
-            valid_joint_start_states = self.mdp.get_valid_joint_player_positions_and_orientations()
+            valid_joint_start_states = (
+                self.mdp.get_valid_joint_player_positions_and_orientations()
+            )
         else:
-            valid_joint_start_states = self.mdp.get_valid_joint_player_positions()
+            valid_joint_start_states = (
+                self.mdp.get_valid_joint_player_positions()
+            )
 
-        valid_player_states = self.mdp.get_valid_player_positions_and_orientations()
-        possible_joint_goal_states = list(itertools.product(valid_player_states, repeat=2))
-        valid_joint_goal_states = list(filter(self.is_valid_joint_motion_goal, possible_joint_goal_states))
+        valid_player_states = (
+            self.mdp.get_valid_player_positions_and_orientations()
+        )
+        possible_joint_goal_states = list(
+            itertools.product(valid_player_states, repeat=2)
+        )
+        valid_joint_goal_states = list(
+            filter(self.is_valid_joint_motion_goal, possible_joint_goal_states)
+        )
 
-        if self.debug: print("Number of plans being pre-calculated: ", len(valid_joint_start_states) * len(valid_joint_goal_states))
-        for joint_start_state, joint_goal_state in itertools.product(valid_joint_start_states, valid_joint_goal_states):
+        if self.debug:
+            print(
+                "Number of plans being pre-calculated: ",
+                len(valid_joint_start_states) * len(valid_joint_goal_states),
+            )
+        for joint_start_state, joint_goal_state in itertools.product(
+            valid_joint_start_states, valid_joint_goal_states
+        ):
             # If orientations not present, joint_start_state just includes positions.
             if not self.start_orientations:
                 dummy_orientation = Direction.NORTH
-                joint_start_state = tuple((pos, dummy_orientation) for pos in joint_start_state)
+                joint_start_state = tuple(
+                    (pos, dummy_orientation) for pos in joint_start_state
+                )
 
             # If either start-end states are not connected, skip to next plan
-            if not self.is_valid_jm_start_goal_pair(joint_start_state, joint_goal_state):
+            if not self.is_valid_jm_start_goal_pair(
+                joint_start_state, joint_goal_state
+            ):
                 continue
 
             # Note: we might fail to get the plan, just due to the nature of the layouts
-            joint_action_list, end_statuses, plan_lengths = self._obtain_plan(joint_start_state, joint_goal_state)
+            joint_action_list, end_statuses, plan_lengths = self._obtain_plan(
+                joint_start_state, joint_goal_state
+            )
             if end_statuses is None:
                 continue
             plan_key = (joint_start_state, joint_goal_state)
-            all_plans[plan_key] = (joint_action_list, end_statuses, plan_lengths)
+            all_plans[plan_key] = (
+                joint_action_list,
+                end_statuses,
+                plan_lengths,
+            )
         return all_plans
 
     def is_valid_jm_start_goal_pair(self, joint_start_state, joint_goal_state):
@@ -457,185 +611,334 @@ def is_valid_jm_start_goal_pair(self, joint_start_state, joint_goal_state):
         if not self.is_valid_joint_motion_goal(joint_goal_state):
             return False
         check_valid_fn = self.motion_planner.is_valid_motion_start_goal_pair
-        return all([check_valid_fn(joint_start_state[i], joint_goal_state[i]) for i in range(2)])
+        return all(
+            [
+                check_valid_fn(joint_start_state[i], joint_goal_state[i])
+                for i in range(2)
+            ]
+        )
 
     def _obtain_plan(self, joint_start_state, joint_goal_state):
         """Either use motion planner or actually compute a joint plan"""
         # Try using MotionPlanner plans and join them together
-        action_plans, pos_and_or_paths, plan_lengths = self._get_plans_from_single_planner(joint_start_state, joint_goal_state)
-        
+        (
+            action_plans,
+            pos_and_or_paths,
+            plan_lengths,
+        ) = self._get_plans_from_single_planner(
+            joint_start_state, joint_goal_state
+        )
+
         # Check if individual plans conflict
-        have_conflict = self.plans_have_conflict(joint_start_state, joint_goal_state, pos_and_or_paths, plan_lengths)
+        have_conflict = self.plans_have_conflict(
+            joint_start_state, joint_goal_state, pos_and_or_paths, plan_lengths
+        )
 
         # If there is no conflict, the joint plan computed by joining single agent MotionPlanner plans is optimal
         if not have_conflict:
-            joint_action_plan, end_pos_and_orientations = self._join_single_agent_action_plans(
-                joint_start_state, action_plans, pos_and_or_paths, min(plan_lengths)
+            (
+                joint_action_plan,
+                end_pos_and_orientations,
+            ) = self._join_single_agent_action_plans(
+                joint_start_state,
+                action_plans,
+                pos_and_or_paths,
+                min(plan_lengths),
             )
             return joint_action_plan, end_pos_and_orientations, plan_lengths
-        
+
         # If there is a conflict in the single motion plan and the agents have the same goal,
         # the graph problem can't be used either as it can't handle same goal state: we compute
         # manually what the best way to handle the conflict is
         elif self._agents_are_in_same_position(joint_goal_state):
-            joint_action_plan, end_pos_and_orientations, plan_lengths = self._handle_path_conflict_with_same_goal(
-                joint_start_state, joint_goal_state, action_plans, pos_and_or_paths
+            (
+                joint_action_plan,
+                end_pos_and_orientations,
+                plan_lengths,
+            ) = self._handle_path_conflict_with_same_goal(
+                joint_start_state,
+                joint_goal_state,
+                action_plans,
+                pos_and_or_paths,
             )
             return joint_action_plan, end_pos_and_orientations, plan_lengths
 
         # If there is a conflict, and the agents have different goals, we can use solve the joint graph problem
-        return self._compute_plan_from_joint_graph(joint_start_state, joint_goal_state)
+        return self._compute_plan_from_joint_graph(
+            joint_start_state, joint_goal_state
+        )
 
-    def _get_plans_from_single_planner(self, joint_start_state, joint_goal_state):
+    def _get_plans_from_single_planner(
+        self, joint_start_state, joint_goal_state
+    ):
         """
         Get individual action plans for each agent from the MotionPlanner to get each agent
         independently to their goal state. NOTE: these plans might conflict
         """
-        single_agent_motion_plans = [self.motion_planner.get_plan(start, goal) for start, goal in zip(joint_start_state, joint_goal_state)]
+        single_agent_motion_plans = [
+            self.motion_planner.get_plan(start, goal)
+            for start, goal in zip(joint_start_state, joint_goal_state)
+        ]
         action_plans, pos_and_or_paths = [], []
         for action_plan, pos_and_or_path, _ in single_agent_motion_plans:
             action_plans.append(action_plan)
             pos_and_or_paths.append(pos_and_or_path)
         plan_lengths = tuple(len(p) for p in action_plans)
-        assert all([plan_lengths[i] == len(pos_and_or_paths[i]) for i in range(2)])
+        assert all(
+            [plan_lengths[i] == len(pos_and_or_paths[i]) for i in range(2)]
+        )
         return action_plans, pos_and_or_paths, plan_lengths
 
-    def plans_have_conflict(self, joint_start_state, joint_goal_state, pos_and_or_paths, plan_lengths):
+    def plans_have_conflict(
+        self,
+        joint_start_state,
+        joint_goal_state,
+        pos_and_or_paths,
+        plan_lengths,
+    ):
         """Check if the sequence of pos_and_or_paths for the two agents conflict"""
         min_length = min(plan_lengths)
         prev_positions = tuple(s[0] for s in joint_start_state)
         for t in range(min_length):
-            curr_pos_or0, curr_pos_or1 = pos_and_or_paths[0][t], pos_and_or_paths[1][t]
+            curr_pos_or0, curr_pos_or1 = (
+                pos_and_or_paths[0][t],
+                pos_and_or_paths[1][t],
+            )
             curr_positions = (curr_pos_or0[0], curr_pos_or1[0])
-            if self.mdp.is_transition_collision(prev_positions, curr_positions):
+            if self.mdp.is_transition_collision(
+                prev_positions, curr_positions
+            ):
                 return True
             prev_positions = curr_positions
         return False
 
-    def _join_single_agent_action_plans(self, joint_start_state, action_plans, pos_and_or_paths, finishing_time):
+    def _join_single_agent_action_plans(
+        self, joint_start_state, action_plans, pos_and_or_paths, finishing_time
+    ):
         """Returns the joint action plan and end joint state obtained by joining the individual action plans"""
         assert finishing_time > 0
-        end_joint_state = (pos_and_or_paths[0][finishing_time - 1], pos_and_or_paths[1][finishing_time - 1])
-        joint_action_plan = list(zip(*[action_plans[0][:finishing_time], action_plans[1][:finishing_time]]))
+        end_joint_state = (
+            pos_and_or_paths[0][finishing_time - 1],
+            pos_and_or_paths[1][finishing_time - 1],
+        )
+        joint_action_plan = list(
+            zip(
+                *[
+                    action_plans[0][:finishing_time],
+                    action_plans[1][:finishing_time],
+                ]
+            )
+        )
         return joint_action_plan, end_joint_state
 
-    def _handle_path_conflict_with_same_goal(self, joint_start_state, joint_goal_state, action_plans, pos_and_or_paths):
-        """Assumes that optimal path in case two agents have the same goal and their paths conflict 
-        is for one of the agents to wait. Checks resulting plans if either agent waits, and selects the 
+    def _handle_path_conflict_with_same_goal(
+        self,
+        joint_start_state,
+        joint_goal_state,
+        action_plans,
+        pos_and_or_paths,
+    ):
+        """Assumes that optimal path in case two agents have the same goal and their paths conflict
+        is for one of the agents to wait. Checks resulting plans if either agent waits, and selects the
         shortest cost among the two."""
 
-        joint_plan0, end_pos_and_or0, plan_lengths0 = self._handle_conflict_with_same_goal_idx(
-            joint_start_state, joint_goal_state, action_plans, pos_and_or_paths, wait_agent_idx=0
+        (
+            joint_plan0,
+            end_pos_and_or0,
+            plan_lengths0,
+        ) = self._handle_conflict_with_same_goal_idx(
+            joint_start_state,
+            joint_goal_state,
+            action_plans,
+            pos_and_or_paths,
+            wait_agent_idx=0,
         )
 
-        joint_plan1, end_pos_and_or1, plan_lengths1 = self._handle_conflict_with_same_goal_idx(
-            joint_start_state, joint_goal_state, action_plans, pos_and_or_paths, wait_agent_idx=1
+        (
+            joint_plan1,
+            end_pos_and_or1,
+            plan_lengths1,
+        ) = self._handle_conflict_with_same_goal_idx(
+            joint_start_state,
+            joint_goal_state,
+            action_plans,
+            pos_and_or_paths,
+            wait_agent_idx=1,
         )
 
         assert any([joint_plan0 is not None, joint_plan1 is not None])
 
         best_plan_idx = np.argmin([min(plan_lengths0), min(plan_lengths1)])
-        solutions = [(joint_plan0, end_pos_and_or0, plan_lengths0), (joint_plan1, end_pos_and_or1, plan_lengths1)]
+        solutions = [
+            (joint_plan0, end_pos_and_or0, plan_lengths0),
+            (joint_plan1, end_pos_and_or1, plan_lengths1),
+        ]
         return solutions[best_plan_idx]
 
-    def _handle_conflict_with_same_goal_idx(self, joint_start_state, joint_goal_state, action_plans, pos_and_or_paths, wait_agent_idx):
+    def _handle_conflict_with_same_goal_idx(
+        self,
+        joint_start_state,
+        joint_goal_state,
+        action_plans,
+        pos_and_or_paths,
+        wait_agent_idx,
+    ):
         """
         Determines what is the best joint plan if whenether there is a conflict between the two agents' actions,
         the agent with index `wait_agent_idx` waits one turn.
-        
+
         If the agent that is assigned to wait is "in front" of the non-waiting agent, this could result
         in an endless conflict. In this case, we return infinite finishing times.
         """
         idx0, idx1 = 0, 0
-        prev_positions = [start_pos_and_or[0] for start_pos_and_or in joint_start_state]
+        prev_positions = [
+            start_pos_and_or[0] for start_pos_and_or in joint_start_state
+        ]
         curr_pos_or0, curr_pos_or1 = joint_start_state
 
         agent0_plan_original, agent1_plan_original = action_plans
 
         joint_plan = []
         # While either agent hasn't finished their plan
-        while idx0 != len(agent0_plan_original) and idx1 != len(agent1_plan_original): 
-            next_pos_or0, next_pos_or1 = pos_and_or_paths[0][idx0], pos_and_or_paths[1][idx1]
+        while idx0 != len(agent0_plan_original) and idx1 != len(
+            agent1_plan_original
+        ):
+            next_pos_or0, next_pos_or1 = (
+                pos_and_or_paths[0][idx0],
+                pos_and_or_paths[1][idx1],
+            )
             next_positions = (next_pos_or0[0], next_pos_or1[0])
 
             # If agents collide, let the waiting agent wait and the non-waiting
             # agent take a step
-            if self.mdp.is_transition_collision(prev_positions, next_positions):
+            if self.mdp.is_transition_collision(
+                prev_positions, next_positions
+            ):
                 if wait_agent_idx == 0:
-                    curr_pos_or0 = curr_pos_or0 # Agent 0 will wait, stays the same
+                    curr_pos_or0 = (
+                        curr_pos_or0  # Agent 0 will wait, stays the same
+                    )
                     curr_pos_or1 = next_pos_or1
-                    curr_joint_action = [Action.STAY, agent1_plan_original[idx1]]
+                    curr_joint_action = [
+                        Action.STAY,
+                        agent1_plan_original[idx1],
+                    ]
                     idx1 += 1
                 elif wait_agent_idx == 1:
                     curr_pos_or0 = next_pos_or0
-                    curr_pos_or1 = curr_pos_or1 # Agent 1 will wait, stays the same
-                    curr_joint_action = [agent0_plan_original[idx0], Action.STAY]
+                    curr_pos_or1 = (
+                        curr_pos_or1  # Agent 1 will wait, stays the same
+                    )
+                    curr_joint_action = [
+                        agent0_plan_original[idx0],
+                        Action.STAY,
+                    ]
                     idx0 += 1
 
                 curr_positions = (curr_pos_or0[0], curr_pos_or1[0])
-                
+
                 # If one agent waiting causes other to crash into it, return None
-                if self._agents_are_in_same_position((curr_pos_or0, curr_pos_or1)):
+                if self._agents_are_in_same_position(
+                    (curr_pos_or0, curr_pos_or1)
+                ):
                     return None, None, [np.Inf, np.Inf]
-                
+
             else:
                 curr_pos_or0, curr_pos_or1 = next_pos_or0, next_pos_or1
                 curr_positions = next_positions
-                curr_joint_action = [agent0_plan_original[idx0], agent1_plan_original[idx1]]
+                curr_joint_action = [
+                    agent0_plan_original[idx0],
+                    agent1_plan_original[idx1],
+                ]
                 idx0 += 1
                 idx1 += 1
-                
+
             joint_plan.append(curr_joint_action)
             prev_positions = curr_positions
-        
+
         assert idx0 != idx1, "No conflict found"
 
         end_pos_and_or = (curr_pos_or0, curr_pos_or1)
-        finishing_times = (np.Inf, idx1) if wait_agent_idx == 0 else (idx0, np.Inf)
+        finishing_times = (
+            (np.Inf, idx1) if wait_agent_idx == 0 else (idx0, np.Inf)
+        )
         return joint_plan, end_pos_and_or, finishing_times
 
     def is_valid_joint_motion_goal(self, joint_goal_state):
         """Checks whether the goal joint positions and orientations are a valid goal"""
-        if not self.same_motion_goals and self._agents_are_in_same_position(joint_goal_state):
+        if not self.same_motion_goals and self._agents_are_in_same_position(
+            joint_goal_state
+        ):
             return False
-        multi_cc_map = len(self.motion_planner.graph_problem.connected_components) > 1
-        players_in_same_cc = self.motion_planner.graph_problem.are_in_same_cc(joint_goal_state[0], joint_goal_state[1])
+        multi_cc_map = (
+            len(self.motion_planner.graph_problem.connected_components) > 1
+        )
+        players_in_same_cc = self.motion_planner.graph_problem.are_in_same_cc(
+            joint_goal_state[0], joint_goal_state[1]
+        )
         if multi_cc_map and players_in_same_cc:
             return False
-        return all([self.motion_planner.is_valid_motion_goal(player_state) for player_state in joint_goal_state])
+        return all(
+            [
+                self.motion_planner.is_valid_motion_goal(player_state)
+                for player_state in joint_goal_state
+            ]
+        )
 
     def is_valid_joint_motion_pair(self, joint_start_state, joint_goal_state):
         if not self.is_valid_joint_motion_goal(joint_goal_state):
             return False
-        return all([ \
-            self.motion_planner.is_valid_motion_start_goal_pair(joint_start_state[i], joint_goal_state[i]) for i in range(2)])
+        return all(
+            [
+                self.motion_planner.is_valid_motion_start_goal_pair(
+                    joint_start_state[i], joint_goal_state[i]
+                )
+                for i in range(2)
+            ]
+        )
 
     def _agents_are_in_same_position(self, joint_motion_state):
-        agent_positions = [player_pos_and_or[0] for player_pos_and_or in joint_motion_state]
+        agent_positions = [
+            player_pos_and_or[0] for player_pos_and_or in joint_motion_state
+        ]
         return len(agent_positions) != len(set(agent_positions))
 
-    def _compute_plan_from_joint_graph(self, joint_start_state, joint_goal_state):
-        """Compute joint action plan for two agents to achieve a 
+    def _compute_plan_from_joint_graph(
+        self, joint_start_state, joint_goal_state
+    ):
+        """Compute joint action plan for two agents to achieve a
         certain position and orientation with the joint motion graph
-        
+
         Args:
             joint_start_state: pair of start (pos, or)
             joint_goal_state: pair of goal (pos, or)
         """
-        assert self.is_valid_joint_motion_pair(joint_start_state, joint_goal_state), joint_goal_state
+        assert self.is_valid_joint_motion_pair(
+            joint_start_state, joint_goal_state
+        ), joint_goal_state
         # Solve shortest-path graph problem
         start_positions = list(zip(*joint_start_state))[0]
         goal_positions = list(zip(*joint_goal_state))[0]
         try:
-            joint_positions_node_path = self.joint_graph_problem.get_node_path(start_positions, goal_positions)[1:]
+            joint_positions_node_path = self.joint_graph_problem.get_node_path(
+                start_positions, goal_positions
+            )[1:]
         except NotConnectedError:
             # The cost will be infinite if there is no path
             num_player = len(goal_positions)
             return [], None, [np.inf] * num_player
-        joint_actions_list, end_pos_and_orientations, finishing_times = self.joint_action_plan_from_positions(joint_positions_node_path, joint_start_state, joint_goal_state)
+        (
+            joint_actions_list,
+            end_pos_and_orientations,
+            finishing_times,
+        ) = self.joint_action_plan_from_positions(
+            joint_positions_node_path, joint_start_state, joint_goal_state
+        )
         return joint_actions_list, end_pos_and_orientations, finishing_times
 
-    def joint_action_plan_from_positions(self, joint_positions, joint_start_state, joint_goal_state):
+    def joint_action_plan_from_positions(
+        self, joint_positions, joint_start_state, joint_goal_state
+    ):
         """
         Finds an action plan and it's cost, such that at least one of the agent goal states is achieved
 
@@ -647,15 +950,22 @@ def joint_action_plan_from_positions(self, joint_positions, joint_start_state, j
         """
         action_plans = []
         for i in range(2):
-            agent_position_sequence = [joint_position[i] for joint_position in joint_positions]
+            agent_position_sequence = [
+                joint_position[i] for joint_position in joint_positions
+            ]
             action_plan, _, _ = self.motion_planner.action_plan_from_positions(
-                                    agent_position_sequence, joint_start_state[i], joint_goal_state[i])
+                agent_position_sequence,
+                joint_start_state[i],
+                joint_goal_state[i],
+            )
             action_plans.append(action_plan)
 
         finishing_times = tuple(len(plan) for plan in action_plans)
         trimmed_action_plans = self._fix_plan_lengths(action_plans)
         joint_action_plan = list(zip(*trimmed_action_plans))
-        end_pos_and_orientations = self._rollout_end_pos_and_or(joint_start_state, joint_action_plan)
+        end_pos_and_orientations = self._rollout_end_pos_and_or(
+            joint_start_state, joint_action_plan
+        )
         return joint_action_plan, end_pos_and_orientations, finishing_times
 
     def _fix_plan_lengths(self, plans):
@@ -665,7 +975,9 @@ def _fix_plan_lengths(self, plans):
         delta_length = max(finishing_times) - min(finishing_times)
         if delta_length != 0:
             index_long_plan = np.argmax(finishing_times)
-            plans[index_long_plan] = plans[index_long_plan][:min(finishing_times)]
+            plans[index_long_plan] = plans[index_long_plan][
+                : min(finishing_times)
+            ]
         return plans
 
     def _rollout_end_pos_and_or(self, joint_start_state, joint_action_plan):
@@ -675,9 +987,16 @@ def _rollout_end_pos_and_or(self, joint_start_state, joint_action_plan):
         # Also assumes can't deliver more than two orders in one motion goal
         # (otherwise Environment will terminate)
         from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
-        dummy_state = OvercookedState.from_players_pos_and_or(joint_start_state, all_orders=self.mdp.start_all_orders)
-        env = OvercookedEnv.from_mdp(self.mdp, horizon=200, info_level=int(self.debug)) # Plans should be shorter than 200 timesteps, or something is likely wrong
-        successor_state, is_done = env.execute_plan(dummy_state, joint_action_plan)
+
+        dummy_state = OvercookedState.from_players_pos_and_or(
+            joint_start_state, all_orders=self.mdp.start_all_orders
+        )
+        env = OvercookedEnv.from_mdp(
+            self.mdp, horizon=200, info_level=int(self.debug)
+        )  # Plans should be shorter than 200 timesteps, or something is likely wrong
+        successor_state, is_done = env.execute_plan(
+            dummy_state, joint_action_plan
+        )
         assert not is_done
         return successor_state.players_pos_and_or
 
@@ -689,25 +1008,36 @@ def _joint_graph_from_grid(self):
         for state_index, joint_pos in enumerate(valid_joint_positions):
             state_decoder[state_index] = joint_pos
 
-        state_encoder = {v:k for k, v in state_decoder.items()}
+        state_encoder = {v: k for k, v in state_decoder.items()}
         num_graph_nodes = len(state_decoder)
 
         adjacency_matrix = np.zeros((num_graph_nodes, num_graph_nodes))
         for start_state_index, start_joint_positions in state_decoder.items():
-            for joint_action, successor_jm_state in self._get_valid_successor_joint_positions(start_joint_positions).items():
+            for (
+                joint_action,
+                successor_jm_state,
+            ) in self._get_valid_successor_joint_positions(
+                start_joint_positions
+            ).items():
                 successor_node_index = state_encoder[successor_jm_state]
 
                 this_action_cost = self._graph_joint_action_cost(joint_action)
-                current_cost = adjacency_matrix[start_state_index][successor_node_index]
+                current_cost = adjacency_matrix[start_state_index][
+                    successor_node_index
+                ]
 
                 if current_cost == 0 or this_action_cost < current_cost:
-                    adjacency_matrix[start_state_index][successor_node_index] = this_action_cost
+                    adjacency_matrix[start_state_index][
+                        successor_node_index
+                    ] = this_action_cost
 
         return Graph(adjacency_matrix, state_encoder, state_decoder)
 
     def _graph_joint_action_cost(self, joint_action):
         """The cost used in the graph shortest-path problem for a certain joint-action"""
-        num_of_non_stay_actions = len([a for a in joint_action if a != Action.STAY])
+        num_of_non_stay_actions = len(
+            [a for a in joint_action if a != Action.STAY]
+        )
         # NOTE: Removing the possibility of having 0 cost joint_actions
         if num_of_non_stay_actions == 0:
             return 1
@@ -717,13 +1047,19 @@ def _get_valid_successor_joint_positions(self, starting_positions):
         """Get all joint positions that can be reached by a joint action.
         NOTE: this DOES NOT include joint positions with superimposed agents."""
         successor_joint_positions = {}
-        joint_motion_actions = itertools.product(Action.MOTION_ACTIONS, Action.MOTION_ACTIONS)
-        
+        joint_motion_actions = itertools.product(
+            Action.MOTION_ACTIONS, Action.MOTION_ACTIONS
+        )
+
         # Under assumption that orientation doesn't matter
         dummy_orientation = Direction.NORTH
-        dummy_player_states = [PlayerState(pos, dummy_orientation) for pos in starting_positions]
+        dummy_player_states = [
+            PlayerState(pos, dummy_orientation) for pos in starting_positions
+        ]
         for joint_action in joint_motion_actions:
-            new_positions, _ = self.mdp.compute_new_positions_and_orientations(dummy_player_states, joint_action)
+            new_positions, _ = self.mdp.compute_new_positions_and_orientations(
+                dummy_player_states, joint_action
+            )
             successor_joint_positions[joint_action] = new_positions
         return successor_joint_positions
 
@@ -742,7 +1078,7 @@ def derive_state(self, start_state, end_pos_and_ors, action_plans):
             position, orientation = end_pos_and_or
             new_player.update_pos_and_or(position, orientation)
             end_players.append(new_player)
-        
+
         end_state.players = tuple(end_players)
 
         # Resolve environment effects for t - 1 turns
@@ -752,9 +1088,14 @@ def derive_state(self, start_state, end_pos_and_ors, action_plans):
             self.mdp.step_environment_effects(end_state)
 
         # Interacts
-        last_joint_action = tuple(a if a == Action.INTERACT else Action.STAY for a in action_plans[-1])
+        last_joint_action = tuple(
+            a if a == Action.INTERACT else Action.STAY
+            for a in action_plans[-1]
+        )
 
-        events_dict = { k : [ [] for _ in range(self.mdp.num_players) ] for k in EVENT_TYPES }
+        events_dict = {
+            k: [[] for _ in range(self.mdp.num_players)] for k in EVENT_TYPES
+        }
         self.mdp.resolve_interacts(end_state, last_joint_action, events_dict)
         self.mdp.resolve_movement(end_state, last_joint_action)
         self.mdp.step_environment_effects(end_state)
@@ -763,9 +1104,9 @@ def derive_state(self, start_state, end_pos_and_ors, action_plans):
 
 class MediumLevelActionManager(object):
     """
-    Manager for medium level actions (specific joint motion goals). 
+    Manager for medium level actions (specific joint motion goals).
     Determines available medium level actions for each state.
-    
+
     Args:
         mdp (OvercookedGridWorld): gridworld of interest
         mlam_params (dictionary): parameters for the medium level action manager
@@ -773,17 +1114,17 @@ class MediumLevelActionManager(object):
 
     def __init__(self, mdp, mlam_params):
         self.mdp = mdp
-        
+
         self.params = mlam_params
-        self.wait_allowed = mlam_params['wait_allowed']
+        self.wait_allowed = mlam_params["wait_allowed"]
         self.counter_drop = mlam_params["counter_drop"]
         self.counter_pickup = mlam_params["counter_pickup"]
-        
+
         self.joint_motion_planner = JointMotionPlanner(mdp, mlam_params)
         self.motion_planner = self.joint_motion_planner.motion_planner
 
     def save_to_file(self, filename):
-        with open(filename, 'wb') as output:
+        with open(filename, "wb") as output:
             pickle.dump(self, output, pickle.HIGHEST_PROTOCOL)
 
     @staticmethod
@@ -791,89 +1132,159 @@ def from_file(filename):
         return load_saved_action_manager(filename)
 
     @staticmethod
-    def from_pickle_or_compute(mdp, mlam_params, custom_filename=None, force_compute=False, info=False):
+    def from_pickle_or_compute(
+        mdp, mlam_params, custom_filename=None, force_compute=False, info=False
+    ):
         assert isinstance(mdp, OvercookedGridworld)
 
-        filename = custom_filename if custom_filename is not None else mdp.layout_name + "_am.pkl"
+        filename = (
+            custom_filename
+            if custom_filename is not None
+            else mdp.layout_name + "_am.pkl"
+        )
 
         if force_compute:
-            return MediumLevelActionManager.compute_mlam(filename, mdp, mlam_params, info=info)
+            return MediumLevelActionManager.compute_mlam(
+                filename, mdp, mlam_params, info=info
+            )
 
         try:
             mlam = MediumLevelActionManager.from_file(filename)
 
             if mlam.params != mlam_params or mlam.mdp != mdp:
                 if info:
-                    print("medium level action manager with different params or mdp found, computing from scratch")
-                return MediumLevelActionManager.compute_mlam(filename, mdp, mlam_params, info=info)
-
-        except (FileNotFoundError, ModuleNotFoundError, EOFError, AttributeError) as e:
+                    print(
+                        "medium level action manager with different params or mdp found, computing from scratch"
+                    )
+                return MediumLevelActionManager.compute_mlam(
+                    filename, mdp, mlam_params, info=info
+                )
+
+        except (
+            FileNotFoundError,
+            ModuleNotFoundError,
+            EOFError,
+            AttributeError,
+        ) as e:
             if info:
                 print("Recomputing planner due to:", e)
-            return MediumLevelActionManager.compute_mlam(filename, mdp, mlam_params, info=info)
+            return MediumLevelActionManager.compute_mlam(
+                filename, mdp, mlam_params, info=info
+            )
 
         if info:
-            print("Loaded MediumLevelActionManager from {}".format(os.path.join(PLANNERS_DIR, filename)))
+            print(
+                "Loaded MediumLevelActionManager from {}".format(
+                    os.path.join(PLANNERS_DIR, filename)
+                )
+            )
         return mlam
 
     @staticmethod
     def compute_mlam(filename, mdp, mlam_params, info=False):
         final_filepath = os.path.join(PLANNERS_DIR, filename)
         if info:
-            print("Computing MediumLevelActionManager to be saved in {}".format(final_filepath))
+            print(
+                "Computing MediumLevelActionManager to be saved in {}".format(
+                    final_filepath
+                )
+            )
         start_time = time.time()
         mlam = MediumLevelActionManager(mdp, mlam_params=mlam_params)
         if info:
-            print("It took {} seconds to create mlam".format(time.time() - start_time))
+            print(
+                "It took {} seconds to create mlam".format(
+                    time.time() - start_time
+                )
+            )
         mlam.save_to_file(final_filepath)
         return mlam
 
     def joint_ml_actions(self, state):
         """Determine all possible joint medium level actions for a certain state"""
-        agent1_actions, agent2_actions = tuple(self.get_medium_level_actions(state, player) for player in state.players)
-        joint_ml_actions = list(itertools.product(agent1_actions, agent2_actions))
-        
+        agent1_actions, agent2_actions = tuple(
+            self.get_medium_level_actions(state, player)
+            for player in state.players
+        )
+        joint_ml_actions = list(
+            itertools.product(agent1_actions, agent2_actions)
+        )
+
         # ml actions are nothing but specific joint motion goals
-        valid_joint_ml_actions = list(filter(lambda a: self.is_valid_ml_action(state, a), joint_ml_actions))
+        valid_joint_ml_actions = list(
+            filter(
+                lambda a: self.is_valid_ml_action(state, a), joint_ml_actions
+            )
+        )
 
         # HACK: Could cause things to break.
         # Necessary to prevent states without successors (due to no counters being allowed and no wait actions)
         # causing A* to not find a solution
         if len(valid_joint_ml_actions) == 0:
-            agent1_actions, agent2_actions = tuple(self.get_medium_level_actions(state, player, waiting_substitute=True) for player in state.players)
-            joint_ml_actions = list(itertools.product(agent1_actions, agent2_actions))
-            valid_joint_ml_actions = list(filter(lambda a: self.is_valid_ml_action(state, a), joint_ml_actions))
+            agent1_actions, agent2_actions = tuple(
+                self.get_medium_level_actions(
+                    state, player, waiting_substitute=True
+                )
+                for player in state.players
+            )
+            joint_ml_actions = list(
+                itertools.product(agent1_actions, agent2_actions)
+            )
+            valid_joint_ml_actions = list(
+                filter(
+                    lambda a: self.is_valid_ml_action(state, a),
+                    joint_ml_actions,
+                )
+            )
             if len(valid_joint_ml_actions) == 0:
-                print("WARNING: Found state without valid actions even after adding waiting substitute actions. State: {}".format(state))
+                print(
+                    "WARNING: Found state without valid actions even after adding waiting substitute actions. State: {}".format(
+                        state
+                    )
+                )
         return valid_joint_ml_actions
 
     def is_valid_ml_action(self, state, ml_action):
-        return self.joint_motion_planner.is_valid_jm_start_goal_pair(state.players_pos_and_or, ml_action)
+        return self.joint_motion_planner.is_valid_jm_start_goal_pair(
+            state.players_pos_and_or, ml_action
+        )
 
-    def get_medium_level_actions(self, state, player, waiting_substitute=False):
+    def get_medium_level_actions(
+        self, state, player, waiting_substitute=False
+    ):
         """
         Determine valid medium level actions for a player.
-        
+
         Args:
             state (OvercookedState): current state
             player (PlayerState): the player's current state
-            waiting_substitute (bool): add a substitute action that takes the place of 
+            waiting_substitute (bool): add a substitute action that takes the place of
                                        a waiting action (going to closest feature)
-        
+
         Returns:
             player_actions (list): possible motion goals (pairs of goal positions and orientations)
         """
         player_actions = []
-        counter_pickup_objects = self.mdp.get_counter_objects_dict(state, self.counter_pickup)
+        counter_pickup_objects = self.mdp.get_counter_objects_dict(
+            state, self.counter_pickup
+        )
         if not player.has_object():
             onion_pickup = self.pickup_onion_actions(counter_pickup_objects)
             tomato_pickup = self.pickup_tomato_actions(counter_pickup_objects)
             dish_pickup = self.pickup_dish_actions(counter_pickup_objects)
-            soup_pickup = self.pickup_counter_soup_actions(counter_pickup_objects)
+            soup_pickup = self.pickup_counter_soup_actions(
+                counter_pickup_objects
+            )
 
             pot_states_dict = self.mdp.get_pot_states(state)
             start_cooking = self.start_cooking_actions(pot_states_dict)
-            player_actions.extend(onion_pickup + tomato_pickup + dish_pickup + soup_pickup + start_cooking)
+            player_actions.extend(
+                onion_pickup
+                + tomato_pickup
+                + dish_pickup
+                + soup_pickup
+                + start_cooking
+            )
 
         else:
             player_object = player.get_object()
@@ -883,17 +1294,25 @@ def get_medium_level_actions(self, state, player, waiting_substitute=False):
             if len(self.counter_drop) > 0:
                 player_actions.extend(self.place_obj_on_counter_actions(state))
 
-            if player_object.name == 'soup':
+            if player_object.name == "soup":
                 player_actions.extend(self.deliver_soup_actions())
-            elif player_object.name == 'onion':
-                player_actions.extend(self.put_onion_in_pot_actions(pot_states_dict))
-            elif player_object.name == 'tomato':
-                player_actions.extend(self.put_tomato_in_pot_actions(pot_states_dict))
-            elif player_object.name == 'dish':
+            elif player_object.name == "onion":
+                player_actions.extend(
+                    self.put_onion_in_pot_actions(pot_states_dict)
+                )
+            elif player_object.name == "tomato":
+                player_actions.extend(
+                    self.put_tomato_in_pot_actions(pot_states_dict)
+                )
+            elif player_object.name == "dish":
                 # Not considering all pots (only ones close to ready) to reduce computation
                 # NOTE: could try to calculate which pots are eligible, but would probably take
                 # a lot of compute
-                player_actions.extend(self.pickup_soup_with_dish_actions(pot_states_dict, only_nearly_ready=False))
+                player_actions.extend(
+                    self.pickup_soup_with_dish_actions(
+                        pot_states_dict, only_nearly_ready=False
+                    )
+                )
             else:
                 raise ValueError("Unrecognized object")
 
@@ -906,42 +1325,53 @@ def get_medium_level_actions(self, state, player, waiting_substitute=False):
             # are not considered valid
             player_actions.extend(self.go_to_closest_feature_actions(player))
 
-        is_valid_goal_given_start = lambda goal: self.motion_planner.is_valid_motion_start_goal_pair(player.pos_and_or, goal)    
-        player_actions = list(filter(is_valid_goal_given_start, player_actions))
+        is_valid_goal_given_start = (
+            lambda goal: self.motion_planner.is_valid_motion_start_goal_pair(
+                player.pos_and_or, goal
+            )
+        )
+        player_actions = list(
+            filter(is_valid_goal_given_start, player_actions)
+        )
         return player_actions
 
     def pickup_onion_actions(self, counter_objects, only_use_dispensers=False):
         """If only_use_dispensers is True, then only take onions from the dispensers"""
         onion_pickup_locations = self.mdp.get_onion_dispenser_locations()
         if not only_use_dispensers:
-            onion_pickup_locations += counter_objects['onion']
+            onion_pickup_locations += counter_objects["onion"]
         return self._get_ml_actions_for_positions(onion_pickup_locations)
 
     def pickup_tomato_actions(self, counter_objects):
         tomato_dispenser_locations = self.mdp.get_tomato_dispenser_locations()
-        tomato_pickup_locations = tomato_dispenser_locations + counter_objects['tomato']
+        tomato_pickup_locations = (
+            tomato_dispenser_locations + counter_objects["tomato"]
+        )
         return self._get_ml_actions_for_positions(tomato_pickup_locations)
 
     def pickup_dish_actions(self, counter_objects, only_use_dispensers=False):
         """If only_use_dispensers is True, then only take dishes from the dispensers"""
         dish_pickup_locations = self.mdp.get_dish_dispenser_locations()
         if not only_use_dispensers:
-            dish_pickup_locations += counter_objects['dish']
+            dish_pickup_locations += counter_objects["dish"]
         return self._get_ml_actions_for_positions(dish_pickup_locations)
 
     def pickup_counter_soup_actions(self, counter_objects):
-        soup_pickup_locations = counter_objects['soup']
+        soup_pickup_locations = counter_objects["soup"]
         return self._get_ml_actions_for_positions(soup_pickup_locations)
 
     def start_cooking_actions(self, pot_states_dict):
         """This is for start cooking a pot that is cookable"""
-        cookable_pots_location = self.mdp.get_partially_full_pots(pot_states_dict) + \
-                                 self.mdp.get_full_but_not_cooking_pots(pot_states_dict)
+        cookable_pots_location = self.mdp.get_partially_full_pots(
+            pot_states_dict
+        ) + self.mdp.get_full_but_not_cooking_pots(pot_states_dict)
         return self._get_ml_actions_for_positions(cookable_pots_location)
 
     def place_obj_on_counter_actions(self, state):
         all_empty_counters = set(self.mdp.get_empty_counter_locations(state))
-        valid_empty_counters = [c_pos for c_pos in self.counter_drop if c_pos in all_empty_counters]
+        valid_empty_counters = [
+            c_pos for c_pos in self.counter_drop if c_pos in all_empty_counters
+        ]
         return self._get_ml_actions_for_positions(valid_empty_counters)
 
     def deliver_soup_actions(self):
@@ -949,39 +1379,67 @@ def deliver_soup_actions(self):
         return self._get_ml_actions_for_positions(serving_locations)
 
     def put_onion_in_pot_actions(self, pot_states_dict):
-        partially_full_onion_pots = self.mdp.get_partially_full_pots(pot_states_dict)
-        fillable_pots = partially_full_onion_pots + pot_states_dict['empty']
+        partially_full_onion_pots = self.mdp.get_partially_full_pots(
+            pot_states_dict
+        )
+        fillable_pots = partially_full_onion_pots + pot_states_dict["empty"]
         return self._get_ml_actions_for_positions(fillable_pots)
 
     def put_tomato_in_pot_actions(self, pot_states_dict):
-        partially_full_onion_pots = self.mdp.get_partially_full_pots(pot_states_dict)
-        fillable_pots = partially_full_onion_pots + pot_states_dict['empty']
+        partially_full_onion_pots = self.mdp.get_partially_full_pots(
+            pot_states_dict
+        )
+        fillable_pots = partially_full_onion_pots + pot_states_dict["empty"]
         return self._get_ml_actions_for_positions(fillable_pots)
 
-    
-    def pickup_soup_with_dish_actions(self, pot_states_dict, only_nearly_ready=False):
-        ready_pot_locations = pot_states_dict['ready']
-        nearly_ready_pot_locations = pot_states_dict['cooking']
+    def pickup_soup_with_dish_actions(
+        self, pot_states_dict, only_nearly_ready=False
+    ):
+        ready_pot_locations = pot_states_dict["ready"]
+        nearly_ready_pot_locations = pot_states_dict["cooking"]
         if not only_nearly_ready:
-            partially_full_pots = self.mdp.get_partially_full_pots(pot_states_dict)
-            nearly_ready_pot_locations = nearly_ready_pot_locations + pot_states_dict['empty'] + partially_full_pots
-        return self._get_ml_actions_for_positions(ready_pot_locations + nearly_ready_pot_locations)
+            partially_full_pots = self.mdp.get_partially_full_pots(
+                pot_states_dict
+            )
+            nearly_ready_pot_locations = (
+                nearly_ready_pot_locations
+                + pot_states_dict["empty"]
+                + partially_full_pots
+            )
+        return self._get_ml_actions_for_positions(
+            ready_pot_locations + nearly_ready_pot_locations
+        )
 
     def go_to_closest_feature_actions(self, player):
-        feature_locations = self.mdp.get_onion_dispenser_locations() + self.mdp.get_tomato_dispenser_locations() + \
-                            self.mdp.get_pot_locations() + self.mdp.get_dish_dispenser_locations()
-        closest_feature_pos = self.motion_planner.min_cost_to_feature(player.pos_and_or, feature_locations, with_argmin=True)[1]
+        feature_locations = (
+            self.mdp.get_onion_dispenser_locations()
+            + self.mdp.get_tomato_dispenser_locations()
+            + self.mdp.get_pot_locations()
+            + self.mdp.get_dish_dispenser_locations()
+        )
+        closest_feature_pos = self.motion_planner.min_cost_to_feature(
+            player.pos_and_or, feature_locations, with_argmin=True
+        )[1]
         return self._get_ml_actions_for_positions([closest_feature_pos])
 
-    def go_to_closest_feature_or_counter_to_goal(self, goal_pos_and_or, goal_location):
+    def go_to_closest_feature_or_counter_to_goal(
+        self, goal_pos_and_or, goal_location
+    ):
         """Instead of going to goal_pos_and_or, go to the closest feature or counter to this goal, that ISN'T the goal itself"""
-        valid_locations = self.mdp.get_onion_dispenser_locations() + \
-                                    self.mdp.get_tomato_dispenser_locations() + self.mdp.get_pot_locations() + \
-                                    self.mdp.get_dish_dispenser_locations() + self.counter_drop
+        valid_locations = (
+            self.mdp.get_onion_dispenser_locations()
+            + self.mdp.get_tomato_dispenser_locations()
+            + self.mdp.get_pot_locations()
+            + self.mdp.get_dish_dispenser_locations()
+            + self.counter_drop
+        )
         valid_locations.remove(goal_location)
         closest_non_goal_feature_pos = self.motion_planner.min_cost_to_feature(
-                                            goal_pos_and_or, valid_locations, with_argmin=True)[1]
-        return self._get_ml_actions_for_positions([closest_non_goal_feature_pos])
+            goal_pos_and_or, valid_locations, with_argmin=True
+        )[1]
+        return self._get_ml_actions_for_positions(
+            [closest_non_goal_feature_pos]
+        )
 
     def wait_actions(self, player):
         waiting_motion_goal = (player.position, player.orientation)
@@ -989,17 +1447,22 @@ def wait_actions(self, player):
 
     def _get_ml_actions_for_positions(self, positions_list):
         """Determine what are the ml actions (joint motion goals) for a list of positions
-        
+
         Args:
             positions_list (list): list of target terrain feature positions
         """
         possible_motion_goals = []
         for pos in positions_list:
             # All possible ways to reach the target feature
-            for motion_goal in self.joint_motion_planner.motion_planner.motion_goals_for_pos[pos]:
+            for (
+                motion_goal
+            ) in self.joint_motion_planner.motion_planner.motion_goals_for_pos[
+                pos
+            ]:
                 possible_motion_goals.append(motion_goal)
         return possible_motion_goals
 
+
 # # Deprecated, since agent-level dynamic planning is no longer used
 # class MediumLevelPlanner(object):
 #     """
@@ -1056,199 +1519,198 @@ def _get_ml_actions_for_positions(self, positions_list):
 #         mlp.ml_action_manager.save_to_file(final_filepath)
 #         return mlp
 #
-    # Deprecated.
-    # def get_successor_states(self, start_state):
-    #     """Successor states for medium-level actions are defined as
-    #     the first state in the corresponding motion plan in which
-    #     one of the two agents' subgoals is satisfied.
-    #
-    #     Returns: list of
-    #         joint_motion_goal: ((pos1, or1), (pos2, or2)) specifying the
-    #                             motion plan goal for both agents
-    #
-    #         successor_state:   OvercookedState corresponding to state
-    #                            arrived at after executing part of the motion plan
-    #                            (until one of the agents arrives at his goal status)
-    #
-    #         plan_length:       Time passed until arrival to the successor state
-    #     """
-    #     if self.mdp.is_terminal(start_state):
-    #         return []
-    #
-    #     start_jm_state = start_state.players_pos_and_or
-    #     successor_states = []
-    #     for goal_jm_state in self.ml_action_manager.joint_ml_actions(start_state):
-    #         joint_motion_action_plans, end_pos_and_ors, plan_costs = self.jmp.get_low_level_action_plan(start_jm_state, goal_jm_state)
-    #         end_state = self.jmp.derive_state(start_state, end_pos_and_ors, joint_motion_action_plans)
-    #
-    #         if SAFE_RUN:
-    #             from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
-    #             assert end_pos_and_ors[0] == goal_jm_state[0] or end_pos_and_ors[1] == goal_jm_state[1]
-    #             s_prime, _ = OvercookedEnv.execute_plan(self.mdp, start_state, joint_motion_action_plans, display=False)
-    #             assert end_state == s_prime,  [self.mdp.state_string(s_prime), self.mdp.state_string(end_state)]
-    #
-    #         successor_states.append((goal_jm_state, end_state, min(plan_costs)))
-    #     return successor_states
-
-    # Deprecated.
-    # def get_successor_states_fixed_other(self, start_state, other_agent, other_agent_idx):
-    #     """
-    #     Get the successor states of a given start state, assuming that the other agent is fixed and will act according to the passed in model
-    #     """
-    #     if self.mdp.is_terminal(start_state):
-    #         return []
-    #
-    #     player = start_state.players[1 - other_agent_idx]
-    #     ml_actions = self.ml_action_manager.get_medium_level_actions(start_state, player)
-    #
-    #     if len(ml_actions) == 0:
-    #         ml_actions = self.ml_action_manager.get_medium_level_actions(start_state, player, waiting_substitute=True)
-    #
-    #     successor_high_level_states = []
-    #     for ml_action in ml_actions:
-    #         action_plan, end_state, cost = self.get_embedded_low_level_action_plan(start_state, ml_action, other_agent, other_agent_idx)
-    #
-    #         if not self.mdp.is_terminal(end_state):
-    #             # Adding interact action and deriving last state
-    #             other_agent_action, _ = other_agent.action(end_state)
-    #             last_joint_action = (Action.INTERACT, other_agent_action) if other_agent_idx == 1 else (other_agent_action, Action.INTERACT)
-    #             action_plan = action_plan + (last_joint_action,)
-    #             cost = cost + 1
-    #
-    #             end_state, _ = self.embedded_mdp_step(end_state, Action.INTERACT, other_agent_action, other_agent.agent_index)
-    #
-    #         successor_high_level_states.append((action_plan, end_state, cost))
-    #     return successor_high_level_states
-
-    # Deprecated. because no longer used
-    # def check_heuristic_consistency(self, curr_heuristic_val, prev_heuristic_val, actual_edge_cost):
-    #     delta_h = curr_heuristic_val - prev_heuristic_val
-    #     assert actual_edge_cost >= delta_h, \
-    #         "Heuristic was not consistent. \n Prev h: {}, Curr h: {}, Actual cost: {}, Δh: {}" \
-    #         .format(prev_heuristic_val, curr_heuristic_val, actual_edge_cost, delta_h)
-    #
-    # def embedded_mdp_succ_fn(self, state, other_agent):
-    #     other_agent_action, _ = other_agent.action(state)
-    #
-    #     successors = []
-    #     for a in Action.ALL_ACTIONS:
-    #         successor_state, joint_action = self.embedded_mdp_step(state, a, other_agent_action, other_agent.agent_index)
-    #         cost = 1
-    #         successors.append((joint_action, successor_state, cost))
-    #     return successors
-    #
-    # def embedded_mdp_step(self, state, action, other_agent_action, other_agent_index):
-    #     if other_agent_index == 0:
-    #         joint_action = (other_agent_action, action)
-    #     else:
-    #         joint_action = (action, other_agent_action)
-    #     if not self.mdp.is_terminal(state):
-    #         results, _ = self.mdp.get_state_transition(state, joint_action)
-    #         successor_state = results
-    #     else:
-    #         print("Tried to find successor of terminal")
-    #         assert False, "state {} \t action {}".format(state, action)
-    #         successor_state = state
-    #     return successor_state, joint_action
-
-    # Deprecated due to Heuristic
-    # def get_low_level_action_plan(self, start_state, h_fn, delivery_horizon=4, debug=False, goal_info=False):
-    #     """
-    #     Get a plan of joint-actions executable in the environment that will lead to a goal number of deliveries
-    #
-    #     Args:
-    #         state (OvercookedState): starting state
-    #         h_fn: heuristic function
-    #
-    #     Returns:
-    #         full_joint_action_plan (list): joint actions to reach goal
-    #     """
-    #     start_state = start_state.deepcopy()
-    #     ml_plan, cost = self.get_ml_plan(start_state, h_fn, delivery_horizon=delivery_horizon, debug=debug)
-    #
-    #     full_joint_action_plan = self.get_low_level_plan_from_ml_plan(
-    #         start_state, ml_plan, h_fn, debug=debug, goal_info=goal_info
-    #     )
-    #     assert cost == len(full_joint_action_plan), "A* cost {} but full joint action plan cost {}".format(cost, len(full_joint_action_plan))
-    #
-    #     if debug: print("Found plan with cost {}".format(cost))
-    #     return full_joint_action_plan
-
-    # Deprecated due to Heuristic
-    # def get_low_level_plan_from_ml_plan(self, start_state, ml_plan, heuristic_fn, debug=False, goal_info=False):
-    #     t = 0
-    #     full_joint_action_plan = []
-    #     curr_state = start_state
-    #     curr_motion_state = start_state.players_pos_and_or
-    #     prev_h = heuristic_fn(start_state, t, debug=False)
-    #
-    #     if len(ml_plan) > 0 and goal_info:
-    #         print("First motion goal: ", ml_plan[0][0])
-    #
-    #     if not clean and debug:
-    #         print("Start state")
-    #         OvercookedEnv.print_state(self.mdp, start_state)
-    #
-    #     for joint_motion_goal, goal_state in ml_plan:
-    #         joint_action_plan, end_motion_state, plan_costs = \
-    #             self.ml_action_manager.joint_motion_planner.get_low_level_action_plan(curr_motion_state, joint_motion_goal)
-    #         curr_plan_cost = min(plan_costs)
-    #         full_joint_action_plan.extend(joint_action_plan)
-    #         t += 1
-    #
-    #         if not clean and debug:
-    #             print(t)
-    #             OvercookedEnv.print_state(self.mdp, goal_state)
-    #
-    #         if not clean and SAFE_RUN:
-    #             s_prime, _ = OvercookedEnv.execute_plan(self.mdp, curr_state, joint_action_plan)
-    #             assert s_prime == goal_state
-    #
-    #         curr_h = heuristic_fn(goal_state, t, debug=False)
-    #         self.check_heuristic_consistency(curr_h, prev_h, curr_plan_cost)
-    #         curr_motion_state, prev_h, curr_state = end_motion_state, curr_h, goal_state
-    #     return full_joint_action_plan
-
-
-    # Deprecated due to Heuristic
-    # def get_ml_plan(self, start_state, h_fn, delivery_horizon=4, debug=False):
-    #     """
-    #     Solves A* Search problem to find optimal sequence of medium level actions
-    #     to reach the goal number of deliveries
-    #
-    #     Returns:
-    #         ml_plan (list): plan not including starting state in form
-    #             [(joint_action, successor_state), ..., (joint_action, goal_state)]
-    #         cost (int): A* Search cost
-    #     """
-    #     start_state = start_state.deepcopy()
-    #
-    #     expand_fn = lambda state: self.get_successor_states(state)
-    #     goal_fn = lambda state: state.delivery_rew >= DELIVERY_REW_THRES
-    #     heuristic_fn = lambda state: h_fn(state)
-    #
-    #     search_problem = SearchTree(start_state, goal_fn, expand_fn, heuristic_fn, debug=debug)
-    #     ml_plan, cost = search_problem.A_star_graph_search(info=True)
-    #     return ml_plan[1:], cost
-
-    # Deprecated
-    # def get_embedded_low_level_action_plan(self, state, goal_pos_and_or, other_agent, other_agent_idx):
-    #     """Find action plan for a specific motion goal with A* considering the other agent"""
-    #     other_agent.set_agent_index(other_agent_idx)
-    #     agent_idx = 1 - other_agent_idx
-    #
-    #     expand_fn = lambda state: self.embedded_mdp_succ_fn(state, other_agent)
-    #     # FIXME
-    #     goal_fn = lambda state: state.players[agent_idx].pos_and_or == goal_pos_and_or or state.delivery_rew >= DELIVERY_REW_THRES
-    #     heuristic_fn = lambda state: sum(pos_distance(state.players[agent_idx].position, goal_pos_and_or[0]))
-    #
-    #     search_problem = SearchTree(state, goal_fn, expand_fn, heuristic_fn)
-    #     state_action_plan, cost = search_problem.A_star_graph_search(info=False)
-    #     action_plan, state_plan = zip(*state_action_plan)
-    #     action_plan = action_plan[1:]
-    #     end_state = state_plan[-1]
-    #     return action_plan, end_state, cost
+# Deprecated.
+# def get_successor_states(self, start_state):
+#     """Successor states for medium-level actions are defined as
+#     the first state in the corresponding motion plan in which
+#     one of the two agents' subgoals is satisfied.
+#
+#     Returns: list of
+#         joint_motion_goal: ((pos1, or1), (pos2, or2)) specifying the
+#                             motion plan goal for both agents
+#
+#         successor_state:   OvercookedState corresponding to state
+#                            arrived at after executing part of the motion plan
+#                            (until one of the agents arrives at his goal status)
+#
+#         plan_length:       Time passed until arrival to the successor state
+#     """
+#     if self.mdp.is_terminal(start_state):
+#         return []
+#
+#     start_jm_state = start_state.players_pos_and_or
+#     successor_states = []
+#     for goal_jm_state in self.ml_action_manager.joint_ml_actions(start_state):
+#         joint_motion_action_plans, end_pos_and_ors, plan_costs = self.jmp.get_low_level_action_plan(start_jm_state, goal_jm_state)
+#         end_state = self.jmp.derive_state(start_state, end_pos_and_ors, joint_motion_action_plans)
+#
+#         if SAFE_RUN:
+#             from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
+#             assert end_pos_and_ors[0] == goal_jm_state[0] or end_pos_and_ors[1] == goal_jm_state[1]
+#             s_prime, _ = OvercookedEnv.execute_plan(self.mdp, start_state, joint_motion_action_plans, display=False)
+#             assert end_state == s_prime,  [self.mdp.state_string(s_prime), self.mdp.state_string(end_state)]
+#
+#         successor_states.append((goal_jm_state, end_state, min(plan_costs)))
+#     return successor_states
 
+# Deprecated.
+# def get_successor_states_fixed_other(self, start_state, other_agent, other_agent_idx):
+#     """
+#     Get the successor states of a given start state, assuming that the other agent is fixed and will act according to the passed in model
+#     """
+#     if self.mdp.is_terminal(start_state):
+#         return []
+#
+#     player = start_state.players[1 - other_agent_idx]
+#     ml_actions = self.ml_action_manager.get_medium_level_actions(start_state, player)
+#
+#     if len(ml_actions) == 0:
+#         ml_actions = self.ml_action_manager.get_medium_level_actions(start_state, player, waiting_substitute=True)
+#
+#     successor_high_level_states = []
+#     for ml_action in ml_actions:
+#         action_plan, end_state, cost = self.get_embedded_low_level_action_plan(start_state, ml_action, other_agent, other_agent_idx)
+#
+#         if not self.mdp.is_terminal(end_state):
+#             # Adding interact action and deriving last state
+#             other_agent_action, _ = other_agent.action(end_state)
+#             last_joint_action = (Action.INTERACT, other_agent_action) if other_agent_idx == 1 else (other_agent_action, Action.INTERACT)
+#             action_plan = action_plan + (last_joint_action,)
+#             cost = cost + 1
+#
+#             end_state, _ = self.embedded_mdp_step(end_state, Action.INTERACT, other_agent_action, other_agent.agent_index)
+#
+#         successor_high_level_states.append((action_plan, end_state, cost))
+#     return successor_high_level_states
+
+# Deprecated. because no longer used
+# def check_heuristic_consistency(self, curr_heuristic_val, prev_heuristic_val, actual_edge_cost):
+#     delta_h = curr_heuristic_val - prev_heuristic_val
+#     assert actual_edge_cost >= delta_h, \
+#         "Heuristic was not consistent. \n Prev h: {}, Curr h: {}, Actual cost: {}, Δh: {}" \
+#         .format(prev_heuristic_val, curr_heuristic_val, actual_edge_cost, delta_h)
+#
+# def embedded_mdp_succ_fn(self, state, other_agent):
+#     other_agent_action, _ = other_agent.action(state)
+#
+#     successors = []
+#     for a in Action.ALL_ACTIONS:
+#         successor_state, joint_action = self.embedded_mdp_step(state, a, other_agent_action, other_agent.agent_index)
+#         cost = 1
+#         successors.append((joint_action, successor_state, cost))
+#     return successors
+#
+# def embedded_mdp_step(self, state, action, other_agent_action, other_agent_index):
+#     if other_agent_index == 0:
+#         joint_action = (other_agent_action, action)
+#     else:
+#         joint_action = (action, other_agent_action)
+#     if not self.mdp.is_terminal(state):
+#         results, _ = self.mdp.get_state_transition(state, joint_action)
+#         successor_state = results
+#     else:
+#         print("Tried to find successor of terminal")
+#         assert False, "state {} \t action {}".format(state, action)
+#         successor_state = state
+#     return successor_state, joint_action
+
+# Deprecated due to Heuristic
+# def get_low_level_action_plan(self, start_state, h_fn, delivery_horizon=4, debug=False, goal_info=False):
+#     """
+#     Get a plan of joint-actions executable in the environment that will lead to a goal number of deliveries
+#
+#     Args:
+#         state (OvercookedState): starting state
+#         h_fn: heuristic function
+#
+#     Returns:
+#         full_joint_action_plan (list): joint actions to reach goal
+#     """
+#     start_state = start_state.deepcopy()
+#     ml_plan, cost = self.get_ml_plan(start_state, h_fn, delivery_horizon=delivery_horizon, debug=debug)
+#
+#     full_joint_action_plan = self.get_low_level_plan_from_ml_plan(
+#         start_state, ml_plan, h_fn, debug=debug, goal_info=goal_info
+#     )
+#     assert cost == len(full_joint_action_plan), "A* cost {} but full joint action plan cost {}".format(cost, len(full_joint_action_plan))
+#
+#     if debug: print("Found plan with cost {}".format(cost))
+#     return full_joint_action_plan
+
+# Deprecated due to Heuristic
+# def get_low_level_plan_from_ml_plan(self, start_state, ml_plan, heuristic_fn, debug=False, goal_info=False):
+#     t = 0
+#     full_joint_action_plan = []
+#     curr_state = start_state
+#     curr_motion_state = start_state.players_pos_and_or
+#     prev_h = heuristic_fn(start_state, t, debug=False)
+#
+#     if len(ml_plan) > 0 and goal_info:
+#         print("First motion goal: ", ml_plan[0][0])
+#
+#     if not clean and debug:
+#         print("Start state")
+#         OvercookedEnv.print_state(self.mdp, start_state)
+#
+#     for joint_motion_goal, goal_state in ml_plan:
+#         joint_action_plan, end_motion_state, plan_costs = \
+#             self.ml_action_manager.joint_motion_planner.get_low_level_action_plan(curr_motion_state, joint_motion_goal)
+#         curr_plan_cost = min(plan_costs)
+#         full_joint_action_plan.extend(joint_action_plan)
+#         t += 1
+#
+#         if not clean and debug:
+#             print(t)
+#             OvercookedEnv.print_state(self.mdp, goal_state)
+#
+#         if not clean and SAFE_RUN:
+#             s_prime, _ = OvercookedEnv.execute_plan(self.mdp, curr_state, joint_action_plan)
+#             assert s_prime == goal_state
+#
+#         curr_h = heuristic_fn(goal_state, t, debug=False)
+#         self.check_heuristic_consistency(curr_h, prev_h, curr_plan_cost)
+#         curr_motion_state, prev_h, curr_state = end_motion_state, curr_h, goal_state
+#     return full_joint_action_plan
+
+
+# Deprecated due to Heuristic
+# def get_ml_plan(self, start_state, h_fn, delivery_horizon=4, debug=False):
+#     """
+#     Solves A* Search problem to find optimal sequence of medium level actions
+#     to reach the goal number of deliveries
+#
+#     Returns:
+#         ml_plan (list): plan not including starting state in form
+#             [(joint_action, successor_state), ..., (joint_action, goal_state)]
+#         cost (int): A* Search cost
+#     """
+#     start_state = start_state.deepcopy()
+#
+#     expand_fn = lambda state: self.get_successor_states(state)
+#     goal_fn = lambda state: state.delivery_rew >= DELIVERY_REW_THRES
+#     heuristic_fn = lambda state: h_fn(state)
+#
+#     search_problem = SearchTree(start_state, goal_fn, expand_fn, heuristic_fn, debug=debug)
+#     ml_plan, cost = search_problem.A_star_graph_search(info=True)
+#     return ml_plan[1:], cost
+
+# Deprecated
+# def get_embedded_low_level_action_plan(self, state, goal_pos_and_or, other_agent, other_agent_idx):
+#     """Find action plan for a specific motion goal with A* considering the other agent"""
+#     other_agent.set_agent_index(other_agent_idx)
+#     agent_idx = 1 - other_agent_idx
+#
+#     expand_fn = lambda state: self.embedded_mdp_succ_fn(state, other_agent)
+#     # FIXME
+#     goal_fn = lambda state: state.players[agent_idx].pos_and_or == goal_pos_and_or or state.delivery_rew >= DELIVERY_REW_THRES
+#     heuristic_fn = lambda state: sum(pos_distance(state.players[agent_idx].position, goal_pos_and_or[0]))
+#
+#     search_problem = SearchTree(state, goal_fn, expand_fn, heuristic_fn)
+#     state_action_plan, cost = search_problem.A_star_graph_search(info=False)
+#     action_plan, state_plan = zip(*state_action_plan)
+#     action_plan = action_plan[1:]
+#     end_state = state_plan[-1]
+#     return action_plan, end_state, cost
 
 
 # Deprecated.
@@ -1717,5 +2179,3 @@ def _get_ml_actions_for_positions(self, positions_list):
 #             print(heuristic_cost, (pot_to_delivery_costs, dish_to_pot_costs, num_pot_to_delivery, items_to_pot_cost))
 #             print(self.mdp.state_string(state))
 #         return heuristic_cost
-
-
diff --git a/src/overcooked_ai_py/planning/search.py b/src/overcooked_ai_py/planning/search.py
index 54d79798..967df5ea 100644
--- a/src/overcooked_ai_py/planning/search.py
+++ b/src/overcooked_ai_py/planning/search.py
@@ -1,12 +1,15 @@
-import heapq, time
+import heapq
+import time
+
 import numpy as np
 import scipy.sparse
 
+
 class SearchTree(object):
     """
     A class to help perform tree searches of various types. Once a goal state is found, returns a list of tuples
     containing (action, state) pairs. This enables to recover the optimal action and state path.
-    
+
     Args:
         root (state): Initial state in our search
         goal_fn (func): Takes in a state and returns whether it is a goal state
@@ -14,7 +17,15 @@ class SearchTree(object):
         heuristic_fn (func): Takes in a state and returns a heuristic value
     """
 
-    def __init__(self, root, goal_fn, expand_fn, heuristic_fn, max_iter_count=10e6, debug=False):
+    def __init__(
+        self,
+        root,
+        goal_fn,
+        expand_fn,
+        heuristic_fn,
+        max_iter_count=10e6,
+        debug=False,
+    ):
         self.debug = debug
         self.root = root
         self.is_goal = goal_fn
@@ -31,7 +42,13 @@ def A_star_graph_search(self, info=False):
         seen = set()
         pq = PriorityQueue()
 
-        root_node = SearchNode(self.root, action=None, parent=None, action_cost=0, debug=self.debug)
+        root_node = SearchNode(
+            self.root,
+            action=None,
+            parent=None,
+            action_cost=0,
+            debug=self.debug,
+        )
         pq.push(root_node, self.estimated_total_cost(root_node))
         while not pq.isEmpty():
             curr_node = pq.pop()
@@ -48,42 +65,63 @@ def A_star_graph_search(self, info=False):
 
             seen.add(curr_state)
             if iter_count > self.max_iter_count:
-                print("Expanded more than the maximum number of allowed states")
+                print(
+                    "Expanded more than the maximum number of allowed states"
+                )
                 raise TimeoutError("Too many states expanded expanded")
 
             if self.is_goal(curr_state):
                 elapsed_time = time.time() - start_time
-                if info: print("Found goal after: \t{:.2f} seconds,   \t{} state expanded ({:.2f} unique) \t ~{:.2f} expansions/s".format(
-                    elapsed_time, iter_count, len(seen)/iter_count, iter_count/elapsed_time))
+                if info:
+                    print(
+                        "Found goal after: \t{:.2f} seconds,   \t{} state expanded ({:.2f} unique) \t ~{:.2f} expansions/s".format(
+                            elapsed_time,
+                            iter_count,
+                            len(seen) / iter_count,
+                            iter_count / elapsed_time,
+                        )
+                    )
                 return curr_node.get_path(), curr_node.backwards_cost
-            
+
             successors = self.expand(curr_state)
 
             for action, child, cost in successors:
-                child_node = SearchNode(child, action, parent=curr_node, action_cost=cost, debug=self.debug)
+                child_node = SearchNode(
+                    child,
+                    action,
+                    parent=curr_node,
+                    action_cost=cost,
+                    debug=self.debug,
+                )
                 pq.push(child_node, self.estimated_total_cost(child_node))
 
-        print("Path for last node expanded: ", [p[0] for p in curr_node.get_path()])
+        print(
+            "Path for last node expanded: ",
+            [p[0] for p in curr_node.get_path()],
+        )
         print("State of last node expanded: ", curr_node.state)
         print("Successors for last node expanded: ", self.expand(curr_state))
-        raise TimeoutError("A* graph search was unable to find any goal state.")
+        raise TimeoutError(
+            "A* graph search was unable to find any goal state."
+        )
 
     def estimated_total_cost(self, node):
         """
         Calculates the estimated total cost of going from node to goal
-        
+
         Args:
             node (SearchNode): node of the state we are interested in
-        
+
         Returns:
             float: h(s) + g(s), where g is the total backwards cost
         """
         return node.backwards_cost + self.heuristic_fn(node.state)
 
+
 class SearchNode(object):
     """
     A helper class that stores a state, action, and parent tuple and enables to restore paths
-    
+
     Args:
         state (any): Game state corresponding to the node
         action (any): Action that brought to the current state
@@ -113,7 +151,7 @@ def __lt__(self, other):
     def get_path(self):
         """
         Returns the path leading from the earliest parent-less node to the current
-        
+
         Returns:
             List of tuples (action, state) where action is the action that led to the state.
             NOTE: The first entry will be (None, start_state).
@@ -125,10 +163,11 @@ def get_path(self):
             node = node.parent
         return path
 
+
 class Graph(object):
     def __init__(self, dense_adjacency_matrix, encoder, decoder, debug=False):
         """
-        Each graph node is distinguishable by a key, encoded by the encoder into 
+        Each graph node is distinguishable by a key, encoded by the encoder into
         a index that corresponds to that node in the adjacency matrix defining the graph.
 
         Arguments:
@@ -136,12 +175,19 @@ def __init__(self, dense_adjacency_matrix, encoder, decoder, debug=False):
             encoder: Dictionary mapping each graph node key to the adj mtx index it corresponds to
             decoder: Dictionary mapping each adj mtx index to a graph node key
         """
-        self.sparse_adjacency_matrix = scipy.sparse.csr_matrix(dense_adjacency_matrix)
+        self.sparse_adjacency_matrix = scipy.sparse.csr_matrix(
+            dense_adjacency_matrix
+        )
         self.distance_matrix = self.shortest_paths(dense_adjacency_matrix)
         self._encoder = encoder
         self._decoder = decoder
         start_time = time.time()
-        if debug: print("Computing shortest paths took {} seconds".format(time.time() - start_time))
+        if debug:
+            print(
+                "Computing shortest paths took {} seconds".format(
+                    time.time() - start_time
+                )
+            )
         self._ccs = None
 
     @property
@@ -157,7 +203,9 @@ def shortest_paths(self, dense_adjacency_matrix):
         Uses scipy's implementation of shortest paths to compute a distance
         matrix between all elements of the graph
         """
-        csgraph = scipy.sparse.csgraph.csgraph_from_dense(dense_adjacency_matrix)
+        csgraph = scipy.sparse.csgraph.csgraph_from_dense(
+            dense_adjacency_matrix
+        )
         return scipy.sparse.csgraph.shortest_path(csgraph)
 
     def dist(self, node1, node2):
@@ -182,15 +230,20 @@ def _get_children(self, node_index):
         """
         assert node_index is not None
         # NOTE: Assuming successor costs are non-zero
-        _, children_indices = self.sparse_adjacency_matrix.getrow(node_index).nonzero()
+        _, children_indices = self.sparse_adjacency_matrix.getrow(
+            node_index
+        ).nonzero()
         return children_indices
-        
+
     def get_node_path(self, start_node, goal_node):
         """
         Given a start node key and a goal node key, returns a list of
         node keys that trace a shortest path from start to goal.
         """
-        start_index, goal_index = self._encoder[start_node], self._encoder[goal_node]
+        start_index, goal_index = (
+            self._encoder[start_node],
+            self._encoder[goal_node],
+        )
         index_path = self._get_node_index_path(start_index, goal_index)
         node_path = [self._decoder[i] for i in index_path]
         return node_path
@@ -225,13 +278,21 @@ def _get_node_index_path(self, start_index, goal_index):
             # X X X X X         X X X X X
             # This is actually an absolutely impossible transition
             # 08/16/2020 update: This has been addressed by catching NotConnectedError upstream
-            raise NotConnectedError("No path could be found from {} to {}".format(self._decoder[start_index], self._decoder[goal_index])
-                                   + "This could be caused by using another layout's planner on this layout")
+            raise NotConnectedError(
+                "No path could be found from {} to {}".format(
+                    self._decoder[start_index], self._decoder[goal_index]
+                )
+                + "This could be caused by using another layout's planner on this layout"
+            )
 
-        return [start_index] + self._get_node_index_path(best_index, goal_index)
+        return [start_index] + self._get_node_index_path(
+            best_index, goal_index
+        )
 
     def _get_connected_components(self):
-        num_ccs, cc_labels = scipy.sparse.csgraph.connected_components(self.sparse_adjacency_matrix)
+        num_ccs, cc_labels = scipy.sparse.csgraph.connected_components(
+            self.sparse_adjacency_matrix
+        )
         connected_components = [set() for _ in range(num_ccs)]
         for node_index, cc_index in enumerate(cc_labels):
             node = self._decoder[node_index]
@@ -239,14 +300,24 @@ def _get_connected_components(self):
         return connected_components
 
     def are_in_same_cc(self, node1, node2):
-        node1_cc_index = [i for i, cc in enumerate(self.connected_components) if node1 in cc]
-        node2_cc_index = [i for i, cc in enumerate(self.connected_components) if node2 in cc]
-        assert len(node1_cc_index) == len(node2_cc_index) == 1, "Node 1 cc: {} \t Node 2 cc: {}".format(node1_cc_index, node2_cc_index)
+        node1_cc_index = [
+            i for i, cc in enumerate(self.connected_components) if node1 in cc
+        ]
+        node2_cc_index = [
+            i for i, cc in enumerate(self.connected_components) if node2 in cc
+        ]
+        assert (
+            len(node1_cc_index) == len(node2_cc_index) == 1
+        ), "Node 1 cc: {} \t Node 2 cc: {}".format(
+            node1_cc_index, node2_cc_index
+        )
         return node1_cc_index[0] == node2_cc_index[0]
 
+
 class NotConnectedError(Exception):
     pass
 
+
 class PriorityQueue:
     """Taken from UC Berkeley's CS188 project utils.
 
@@ -258,7 +329,8 @@ class PriorityQueue:
     Note that this PriorityQueue does not allow you to change the priority
     of an item. However, you may insert the same item multiple times with
     different priorities."""
-    def  __init__(self):
+
+    def __init__(self):
         self.heap = []
 
     def push(self, item, priority):
diff --git a/src/overcooked_ai_py/static.py b/src/overcooked_ai_py/static.py
index d6d923cc..ab8ffd88 100644
--- a/src/overcooked_ai_py/static.py
+++ b/src/overcooked_ai_py/static.py
@@ -7,4 +7,4 @@
 LAYOUTS_DIR = os.path.join(DATA_DIR, "layouts")
 GRAPHICS_DIR = os.path.join(DATA_DIR, "graphics")
 FONTS_DIR = os.path.join(DATA_DIR, "fonts")
-TESTING_DATA_DIR = os.path.join(DATA_DIR, "testing")
\ No newline at end of file
+TESTING_DATA_DIR = os.path.join(DATA_DIR, "testing")
diff --git a/src/overcooked_ai_py/utils.py b/src/overcooked_ai_py/utils.py
index 1fe467f9..d5cbb7c9 100644
--- a/src/overcooked_ai_py/utils.py
+++ b/src/overcooked_ai_py/utils.py
@@ -1,54 +1,74 @@
-import io, json, pickle, pstats, cProfile, os, tempfile, uuid
-import numpy as np
-from numpy import nan
+import cProfile
+import io
+import json
+import os
+import pickle
+import pstats
+import tempfile
+import uuid
 from collections import defaultdict
-from pathlib import Path
 from collections.abc import Iterable
+from pathlib import Path
+
+import numpy as np
+from numpy import nan
+
 from overcooked_ai_py.static import LAYOUTS_DIR
-from collections.abc import Iterable
 
 # I/O
 
+
 def save_pickle(data, filename):
     with open(fix_filetype(filename, ".pickle"), "wb") as f:
         pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
 
+
 def load_pickle(filename):
     with open(fix_filetype(filename, ".pickle"), "rb") as f:
         return pickle.load(f)
 
+
 def load_dict_from_file(filepath):
     with open(filepath, "r") as f:
         return eval(f.read())
 
+
 def save_dict_to_file(dic, filename):
     dic = dict(dic)
-    with open(fix_filetype(filename, ".txt"),"w") as f:
+    with open(fix_filetype(filename, ".txt"), "w") as f:
         f.write(str(dic))
 
+
 def load_dict_from_txt(filename):
     return load_dict_from_file(fix_filetype(filename, ".txt"))
 
+
 def save_as_json(data, filename):
     with open(fix_filetype(filename, ".json"), "w") as outfile:
         json.dump(data, outfile)
     return filename
 
+
 def load_from_json(filename):
     with open(fix_filetype(filename, ".json"), "r") as json_file:
         return json.load(json_file)
 
+
 def iterate_over_json_files_in_dir(dir_path):
     pathlist = Path(dir_path).glob("*.json")
     return [str(path) for path in pathlist]
 
+
 def fix_filetype(path, filetype):
-    if path[-len(filetype):] == filetype:
+    if path[-len(filetype) :] == filetype:
         return path
     else:
         return path + filetype
 
-def generate_temporary_file_path(file_name=None, prefix="", suffix="", extension=""):
+
+def generate_temporary_file_path(
+    file_name=None, prefix="", suffix="", extension=""
+):
     if file_name is None:
         file_name = str(uuid.uuid1())
     if extension and not extension.startswith("."):
@@ -56,47 +76,60 @@ def generate_temporary_file_path(file_name=None, prefix="", suffix="", extension
     file_name = prefix + file_name + suffix + extension
     return os.path.join(tempfile.gettempdir(), file_name)
 
+
 # MDP
 
+
 def cumulative_rewards_from_rew_list(rews):
     return [sum(rews[:t]) for t in range(len(rews))]
 
+
 # Gridworld
 
+
 def manhattan_distance(pos1, pos2):
     """Returns manhattan distance between two points in (x, y) format"""
     return abs(pos1[0] - pos2[0]) + abs(pos1[1] - pos2[1])
 
+
 def pos_distance(pos0, pos1):
     return tuple(np.array(pos0) - np.array(pos1))
 
+
 # Randomness
 
+
 def rnd_uniform(low, high):
     if low == high:
         return low
     return np.random.uniform(low, high)
 
+
 def rnd_int_uniform(low, high):
     if low == high:
         return low
     return np.random.choice(range(low, high + 1))
 
+
 # Statistics
 
+
 def std_err(lst):
     """Computes the standard error"""
     sd = np.std(lst)
     n = len(lst)
     return sd / np.sqrt(n)
 
+
 def mean_and_std_err(lst):
     "Mean and standard error of list"
     mu = np.mean(lst)
     return mu, std_err(lst)
 
+
 # Other utils
 
+
 def dict_mean_and_std_err(d):
     """
     Takes in a dictionary with lists as keys, and returns a dictionary
@@ -108,35 +141,42 @@ def dict_mean_and_std_err(d):
         result[k] = mean_and_std_err(v)
     return result
 
+
 def append_dictionaries(dictionaries):
     """
     Append many dictionaries with numbers as values into one dictionary with lists as values.
 
     {a: 1, b: 2}, {a: 3, b: 0}  ->  {a: [1, 3], b: [2, 0]}
     """
-    assert all(set(d.keys()) == set(dictionaries[0].keys()) for d in dictionaries), "All key sets are the same across all dicts"
+    assert all(
+        set(d.keys()) == set(dictionaries[0].keys()) for d in dictionaries
+    ), "All key sets are the same across all dicts"
     final_dict = defaultdict(list)
     for d in dictionaries:
         for k, v in d.items():
             final_dict[k].append(v)
     return dict(final_dict)
 
+
 def merge_dictionaries(dictionaries):
     """
     Merge many dictionaries by extending them to one another.
     {a: [1, 7], b: [2, 5]}, {a: [3], b: [0]}  ->  {a: [1, 7, 3], b: [2, 5, 0]}
     """
-    assert all(set(d.keys()) == set(dictionaries[0].keys()) for d in dictionaries), "All key sets are the same across all dicts"
+    assert all(
+        set(d.keys()) == set(dictionaries[0].keys()) for d in dictionaries
+    ), "All key sets are the same across all dicts"
     final_dict = defaultdict(list)
     for d in dictionaries:
         for k, v in d.items():
             final_dict[k].extend(v)
     return dict(final_dict)
 
+
 def rm_idx_from_dict(d, idx):
     """
     Takes in a dictionary with lists as values, and returns
-    a dictionary with lists as values, but containing 
+    a dictionary with lists as values, but containing
     only the desired index
 
     NOTE: this is a MUTATING METHOD, returns the POPPED IDX
@@ -147,21 +187,25 @@ def rm_idx_from_dict(d, idx):
         new_d[k] = [d[k].pop(idx)]
     return new_d
 
+
 def take_indexes_from_dict(d, indices, keys_to_ignore=[]):
     """
     Takes in a dictionary with lists as values, and returns
     a dictionary with lists as values, but with subsampled indices
-    based on the `indices` input 
+    based on the `indices` input
     """
     assert all(isinstance(v, Iterable) for v in d.values())
     new_d = {}
     for k, v in d.items():
-        if k in keys_to_ignore: continue
+        if k in keys_to_ignore:
+            continue
         new_d[k] = np.take(d[k], indices)
     return new_d
 
+
 def profile(fnc):
     """A decorator that uses cProfile to profile a function (from https://osf.io/upav8/)"""
+
     def inner(*args, **kwargs):
         pr = cProfile.Profile()
         pr.enable()
@@ -172,17 +216,24 @@ def inner(*args, **kwargs):
         ps.print_stats()
         print(s.getvalue())
         return retval
+
     return inner
 
+
 def read_layout_dict(layout_name):
-    return load_dict_from_file(os.path.join(LAYOUTS_DIR, layout_name + ".layout"))
+    return load_dict_from_file(
+        os.path.join(LAYOUTS_DIR, layout_name + ".layout")
+    )
+
 
 class classproperty(property):
     def __get__(self, cls, owner):
         return classmethod(self.fget).__get__(None, owner)()
 
+
 class OvercookedException(Exception):
     pass
 
+
 def is_iterable(obj):
     return isinstance(obj, Iterable)
diff --git a/src/overcooked_ai_py/visualization/pygame_utils.py b/src/overcooked_ai_py/visualization/pygame_utils.py
index d0cd6729..a5313805 100644
--- a/src/overcooked_ai_py/visualization/pygame_utils.py
+++ b/src/overcooked_ai_py/visualization/pygame_utils.py
@@ -1,39 +1,48 @@
 import pygame
-from pygame.locals import HWSURFACE, DOUBLEBUF, RESIZABLE, QUIT, VIDEORESIZE
+from pygame.locals import DOUBLEBUF, HWSURFACE, QUIT, RESIZABLE, VIDEORESIZE
+
 from overcooked_ai_py.utils import load_from_json
 
+
 def run_static_resizeable_window(surface, fps=30):
     """
     window that can be resized and closed using gui
     """
     pygame.init()
     clock = pygame.time.Clock()
-    window = pygame.display.set_mode(surface.get_size(), HWSURFACE|DOUBLEBUF|RESIZABLE)
+    window = pygame.display.set_mode(
+        surface.get_size(), HWSURFACE | DOUBLEBUF | RESIZABLE
+    )
     window.blit(surface, (0, 0))
     pygame.display.flip()
     try:
         while True:
             pygame.event.pump()
             event = pygame.event.wait()
-            if event.type == QUIT: 
+            if event.type == QUIT:
                 pygame.display.quit()
                 pygame.quit()
             elif event.type == VIDEORESIZE:
-                window = pygame.display.set_mode(event.dict['size'], HWSURFACE|DOUBLEBUF|RESIZABLE)
-                window.blit(pygame.transform.scale(surface, event.dict['size']), (0, 0))
+                window = pygame.display.set_mode(
+                    event.dict["size"], HWSURFACE | DOUBLEBUF | RESIZABLE
+                )
+                window.blit(
+                    pygame.transform.scale(surface, event.dict["size"]), (0, 0)
+                )
                 pygame.display.flip()
                 clock.tick(fps)
     except:
         pygame.display.quit()
         pygame.quit()
-        if event.type != QUIT: # if user meant to quit error does not matter
+        if event.type != QUIT:  # if user meant to quit error does not matter
             raise
 
+
 def vstack_surfaces(surfaces, background_color=None):
-    '''
+    """
     stack surfaces vertically (on y axis)
     if surfaces have different width fill remaining area with background color
-    '''
+    """
     result_width = max(surface.get_width() for surface in surfaces)
     result_height = sum(surface.get_height() for surface in surfaces)
     result_surface = pygame.surface.Surface((result_width, result_height))
@@ -45,42 +54,57 @@ def vstack_surfaces(surfaces, background_color=None):
         next_surface_y_position += surface.get_height()
     return result_surface
 
+
 def scale_surface_by_factor(surface, scale_by_factor):
-    ''' return scaled input surfacem (with size multiplied by scale_by_factor param)
-        scales also content of the surface
-    '''
+    """return scaled input surfacem (with size multiplied by scale_by_factor param)
+    scales also content of the surface
+    """
     unscaled_size = surface.get_size()
-    scaled_size = tuple(int(dim*scale_by_factor) for dim in unscaled_size)
-    return  pygame.transform.scale(surface, scaled_size)
+    scaled_size = tuple(int(dim * scale_by_factor) for dim in unscaled_size)
+    return pygame.transform.scale(surface, scaled_size)
+
 
 def blit_on_new_surface_of_size(surface, size, background_color=None):
-    '''blit surface on new surface of given size of surface (with no resize of its content), filling not covered parts of result area with background color
-    '''
+    """blit surface on new surface of given size of surface (with no resize of its content), filling not covered parts of result area with background color"""
     result_surface = pygame.surface.Surface(size)
     if background_color:
         result_surface.fill(background_color)
-    result_surface.blit(surface, (0,0))
+    result_surface.blit(surface, (0, 0))
     return result_surface
 
+
 class MultiFramePygameImage:
-    ''' use to read frames of images from overcooked-demo repo easly
-    '''
+    """use to read frames of images from overcooked-demo repo easly"""
+
     def __init__(self, img_path, frames_path):
         self.image = pygame.image.load(img_path)
-        self.frames_rectangles = MultiFramePygameImage.load_frames_rectangles(frames_path)
-    
-    def blit_on_surface(self, surface, top_left_pixel_position, frame_name, **kwargs):
-            surface.blit(self.image, top_left_pixel_position, area=self.frames_rectangles[frame_name], **kwargs)
-    
+        self.frames_rectangles = MultiFramePygameImage.load_frames_rectangles(
+            frames_path
+        )
+
+    def blit_on_surface(
+        self, surface, top_left_pixel_position, frame_name, **kwargs
+    ):
+        surface.blit(
+            self.image,
+            top_left_pixel_position,
+            area=self.frames_rectangles[frame_name],
+            **kwargs
+        )
+
     @staticmethod
     def load_frames_rectangles(json_path):
         frames_json = load_from_json(json_path)
 
-        if "textures" in frames_json.keys(): # check if its format of soups.json
-            assert frames_json["textures"][0]["scale"] == 1 # not implemented support for scale here
+        if (
+            "textures" in frames_json.keys()
+        ):  # check if its format of soups.json
+            assert (
+                frames_json["textures"][0]["scale"] == 1
+            )  # not implemented support for scale here
             frames = frames_json["textures"][0]["frames"]
 
-        else: # assume its format of objects.json, terrain.json and chefs.json
+        else:  # assume its format of objects.json, terrain.json and chefs.json
             frames = []
             for filename, frame_dict in frames_json["frames"].items():
                 frame_dict["filename"] = filename
@@ -88,10 +112,10 @@ def load_frames_rectangles(json_path):
 
         result = {}
         for frame_dict in frames:
-            assert not frame_dict.get("rotated") # not implemented support yet
-            assert not frame_dict.get("trimmed") # not implemented support yet
+            assert not frame_dict.get("rotated")  # not implemented support yet
+            assert not frame_dict.get("trimmed")  # not implemented support yet
             frame_name = frame_dict["filename"].split(".")[0]
             frame = frame_dict["frame"]
             rect = pygame.Rect(frame["x"], frame["y"], frame["w"], frame["h"])
             result[frame_name] = rect
-        return result
\ No newline at end of file
+        return result
diff --git a/src/overcooked_ai_py/visualization/state_visualizer.py b/src/overcooked_ai_py/visualization/state_visualizer.py
index 74c5dcbb..34e5cd0e 100644
--- a/src/overcooked_ai_py/visualization/state_visualizer.py
+++ b/src/overcooked_ai_py/visualization/state_visualizer.py
@@ -1,36 +1,81 @@
+import copy
+import math
+import os
+
 import pygame
-import os, copy, math
-from overcooked_ai_py.utils import generate_temporary_file_path, classproperty, cumulative_rewards_from_rew_list
-from overcooked_ai_py.static import GRAPHICS_DIR, FONTS_DIR
-from overcooked_ai_py.mdp.layout_generator import EMPTY, COUNTER, ONION_DISPENSER, TOMATO_DISPENSER, POT, DISH_DISPENSER, SERVING_LOC
-from overcooked_ai_py.visualization.visualization_utils import show_image_in_ipython, show_ipython_images_slider
-from overcooked_ai_py.visualization.pygame_utils import MultiFramePygameImage, run_static_resizeable_window, vstack_surfaces, scale_surface_by_factor, blit_on_new_surface_of_size
-from overcooked_ai_py.mdp.actions import Direction, Action
+
+from overcooked_ai_py.mdp.actions import Action, Direction
+from overcooked_ai_py.mdp.layout_generator import (
+    COUNTER,
+    DISH_DISPENSER,
+    EMPTY,
+    ONION_DISPENSER,
+    POT,
+    SERVING_LOC,
+    TOMATO_DISPENSER,
+)
+from overcooked_ai_py.static import FONTS_DIR, GRAPHICS_DIR
+from overcooked_ai_py.utils import (
+    classproperty,
+    cumulative_rewards_from_rew_list,
+    generate_temporary_file_path,
+)
+from overcooked_ai_py.visualization.pygame_utils import (
+    MultiFramePygameImage,
+    blit_on_new_surface_of_size,
+    run_static_resizeable_window,
+    scale_surface_by_factor,
+    vstack_surfaces,
+)
+from overcooked_ai_py.visualization.visualization_utils import (
+    show_image_in_ipython,
+    show_ipython_images_slider,
+)
 
 roboto_path = os.path.join(FONTS_DIR, "Roboto-Regular.ttf")
 
 
 class StateVisualizer:
-    TERRAINS_IMG = MultiFramePygameImage(os.path.join(GRAPHICS_DIR, 'terrain.png'), os.path.join(GRAPHICS_DIR, 'terrain.json'))
-    OBJECTS_IMG = MultiFramePygameImage(os.path.join(GRAPHICS_DIR, 'objects.png'), os.path.join(GRAPHICS_DIR, 'objects.json'))
-    SOUPS_IMG = MultiFramePygameImage(os.path.join(GRAPHICS_DIR, 'soups.png'), os.path.join(GRAPHICS_DIR, 'soups.json'))
-    CHEFS_IMG = MultiFramePygameImage(os.path.join(GRAPHICS_DIR, 'chefs.png'), os.path.join(GRAPHICS_DIR, 'chefs.json'))
-    ARROW_IMG = pygame.image.load(os.path.join(GRAPHICS_DIR, 'arrow.png'))
-    INTERACT_IMG = pygame.image.load(os.path.join(GRAPHICS_DIR, 'interact.png'))
-    STAY_IMG = pygame.image.load(os.path.join(GRAPHICS_DIR, 'stay.png'))
+    TERRAINS_IMG = MultiFramePygameImage(
+        os.path.join(GRAPHICS_DIR, "terrain.png"),
+        os.path.join(GRAPHICS_DIR, "terrain.json"),
+    )
+    OBJECTS_IMG = MultiFramePygameImage(
+        os.path.join(GRAPHICS_DIR, "objects.png"),
+        os.path.join(GRAPHICS_DIR, "objects.json"),
+    )
+    SOUPS_IMG = MultiFramePygameImage(
+        os.path.join(GRAPHICS_DIR, "soups.png"),
+        os.path.join(GRAPHICS_DIR, "soups.json"),
+    )
+    CHEFS_IMG = MultiFramePygameImage(
+        os.path.join(GRAPHICS_DIR, "chefs.png"),
+        os.path.join(GRAPHICS_DIR, "chefs.json"),
+    )
+    ARROW_IMG = pygame.image.load(os.path.join(GRAPHICS_DIR, "arrow.png"))
+    INTERACT_IMG = pygame.image.load(
+        os.path.join(GRAPHICS_DIR, "interact.png")
+    )
+    STAY_IMG = pygame.image.load(os.path.join(GRAPHICS_DIR, "stay.png"))
     UNSCALED_TILE_SIZE = 15
     DEFAULT_VALUES = {
-        "height": None, # if None use grid_width - NOTE: can chop down hud if hud is wider than grid
-        "width": None, # if None use (hud_height+grid_height)
+        "height": None,  # if None use grid_width - NOTE: can chop down hud if hud is wider than grid
+        "width": None,  # if None use (hud_height+grid_height)
         "tile_size": 75,
         "window_fps": 30,
-        "player_colors": ['blue', 'green'],
+        "player_colors": ["blue", "green"],
         "is_rendering_hud": True,
         "hud_font_size": 10,
-        "hud_font_path": roboto_path, 
-        "hud_system_font_name": None, # if set to None use hud_font_path
-        "hud_font_color": (255,255,255), # white
-        "hud_data_default_key_order" : ["all_orders", "bonus_orders", "time_left", "score", "potential"],
+        "hud_font_path": roboto_path,
+        "hud_system_font_name": None,  # if set to None use hud_font_path
+        "hud_font_color": (255, 255, 255),  # white
+        "hud_data_default_key_order": [
+            "all_orders",
+            "bonus_orders",
+            "time_left",
+            "score",
+            "potential",
+        ],
         "hud_interline_size": 10,
         "hud_margin_bottom": 10,
         "hud_margin_top": 10,
@@ -39,13 +84,13 @@ class StateVisualizer:
         "hud_order_size": 15,
         "is_rendering_cooking_timer": True,
         "show_timer_when_cooked": True,
-        "cooking_timer_font_size": 20, # # if set to None use cooking_timer_font_path
-        "cooking_timer_font_path": roboto_path, 
+        "cooking_timer_font_size": 20,  # # if set to None use cooking_timer_font_path
+        "cooking_timer_font_path": roboto_path,
         "cooking_timer_system_font_name": None,
-        "cooking_timer_font_color": (255, 0, 0), # red
+        "cooking_timer_font_color": (255, 0, 0),  # red
         "grid": None,
-        "background_color": (155, 101, 0), # color of empty counter
-        "is_rendering_action_probs": True # whatever represent visually on the grid what actions some given agent would make
+        "background_color": (155, 101, 0),  # color of empty counter
+        "is_rendering_action_probs": True,  # whatever represent visually on the grid what actions some given agent would make
     }
     TILE_TO_FRAME_NAME = {
         EMPTY: "floor",
@@ -54,7 +99,7 @@ class StateVisualizer:
         TOMATO_DISPENSER: "tomatoes",
         POT: "pot",
         DISH_DISPENSER: "dishes",
-        SERVING_LOC: "serve"
+        SERVING_LOC: "serve",
     }
 
     def __init__(self, **kwargs):
@@ -67,16 +112,23 @@ def reload_fonts(self):
         pygame.font.init()
         if not hasattr(self, "_font"):
             self._fonts = {}
-        # initializing fonts only if needed because it can take a quite long time, 
+        # initializing fonts only if needed because it can take a quite long time,
         #   see https://pygame.readthedocs.io/en/latest/4_text/text.html#initialize-a-font
         if self.is_rendering_hud:
-            self.hud_font = self._init_font(self.hud_font_size, self.hud_font_path, self.hud_system_font_name)
+            self.hud_font = self._init_font(
+                self.hud_font_size,
+                self.hud_font_path,
+                self.hud_system_font_name,
+            )
         else:
             self.hud_font = None
 
         if self.is_rendering_cooking_timer:
-            self.cooking_timer_font = self._init_font(self.cooking_timer_font_size, self.cooking_timer_font_path,
-                                                     self.cooking_timer_system_font_name)
+            self.cooking_timer_font = self._init_font(
+                self.cooking_timer_font_size,
+                self.cooking_timer_font_path,
+                self.cooking_timer_system_font_name,
+            )
         else:
             self.cooking_timer_font = None
 
@@ -92,17 +144,37 @@ def configure(self, **kwargs):
 
     @staticmethod
     def default_hud_data(state, **kwargs):
-        result = {"timestep": state.timestep, "all_orders":[r.to_dict() for r in state.all_orders], "bonus_orders": [r.to_dict() for r in state.bonus_orders]}
+        result = {
+            "timestep": state.timestep,
+            "all_orders": [r.to_dict() for r in state.all_orders],
+            "bonus_orders": [r.to_dict() for r in state.bonus_orders],
+        }
         result.update(copy.deepcopy(kwargs))
         return result
 
     @staticmethod
     def default_hud_data_from_trajectories(trajectories, trajectory_idx=0):
-        scores = cumulative_rewards_from_rew_list(trajectories["ep_rewards"][trajectory_idx])
-        return [StateVisualizer.default_hud_data(state, score=scores[i])
-            for i, state in enumerate(trajectories["ep_states"][trajectory_idx])]
-
-    def display_rendered_trajectory(self, trajectories, trajectory_idx=0,  hud_data=None, action_probs=None, img_directory_path=None, img_extension=".png", img_prefix="", ipython_display=True):
+        scores = cumulative_rewards_from_rew_list(
+            trajectories["ep_rewards"][trajectory_idx]
+        )
+        return [
+            StateVisualizer.default_hud_data(state, score=scores[i])
+            for i, state in enumerate(
+                trajectories["ep_states"][trajectory_idx]
+            )
+        ]
+
+    def display_rendered_trajectory(
+        self,
+        trajectories,
+        trajectory_idx=0,
+        hud_data=None,
+        action_probs=None,
+        img_directory_path=None,
+        img_extension=".png",
+        img_prefix="",
+        ipython_display=True,
+    ):
         """
         saves images of every timestep from trajectory in img_directory_path (or temporary directory if not path is not specified)
         trajectories (dict): trajectories dict, same format as used by AgentEvaluator
@@ -116,7 +188,9 @@ def display_rendered_trajectory(self, trajectories, trajectory_idx=0,  hud_data=
         grid = trajectories["mdp_params"][trajectory_idx]["terrain"]
         if hud_data is None:
             if self.is_rendering_hud:
-                hud_data = StateVisualizer.default_hud_data_from_trajectories(trajectories, trajectory_idx)
+                hud_data = StateVisualizer.default_hud_data_from_trajectories(
+                    trajectories, trajectory_idx
+                )
             else:
                 hud_data = [None] * len(states)
 
@@ -124,20 +198,41 @@ def display_rendered_trajectory(self, trajectories, trajectory_idx=0,  hud_data=
             action_probs = [None] * len(states)
 
         if not img_directory_path:
-            img_directory_path = generate_temporary_file_path(prefix="overcooked_visualized_trajectory", extension="")
+            img_directory_path = generate_temporary_file_path(
+                prefix="overcooked_visualized_trajectory", extension=""
+            )
         os.makedirs(img_directory_path, exist_ok=True)
         img_pathes = []
         for i, state in enumerate(states):
             img_name = img_prefix + str(i) + img_extension
             img_path = os.path.join(img_directory_path, img_name)
-            img_pathes.append(self.display_rendered_state(state=state, hud_data=hud_data[i], action_probs=action_probs[i], grid=grid, img_path=img_path, ipython_display=False, window_display=False))
+            img_pathes.append(
+                self.display_rendered_state(
+                    state=state,
+                    hud_data=hud_data[i],
+                    action_probs=action_probs[i],
+                    grid=grid,
+                    img_path=img_path,
+                    ipython_display=False,
+                    window_display=False,
+                )
+            )
 
         if ipython_display:
             return show_ipython_images_slider(img_pathes, "timestep")
 
         return img_directory_path
 
-    def display_rendered_state(self, state, hud_data=None, action_probs=None, grid=None, img_path=None, ipython_display=False, window_display=False):
+    def display_rendered_state(
+        self,
+        state,
+        hud_data=None,
+        action_probs=None,
+        grid=None,
+        img_path=None,
+        ipython_display=False,
+        window_display=False,
+    ):
         """
         renders state as image
         state (OvercookedState): state to render
@@ -148,11 +243,17 @@ def display_rendered_state(self, state, hud_data=None, action_probs=None, grid=N
         window_display (bool): if True render state into pygame window
         action_probs(list(list(float))): action probs for every player acessed in the way action_probs[player][action]
         """
-        assert window_display or img_path or ipython_display, "specify at least one of the ways to output result state image: window_display, img_path, or ipython_display"
-        surface = self.render_state(state, grid, hud_data, action_probs=action_probs)
+        assert (
+            window_display or img_path or ipython_display
+        ), "specify at least one of the ways to output result state image: window_display, img_path, or ipython_display"
+        surface = self.render_state(
+            state, grid, hud_data, action_probs=action_probs
+        )
 
         if img_path is None and ipython_display:
-            img_path = generate_temporary_file_path(prefix="overcooked_visualized_state_", extension=".png")
+            img_path = generate_temporary_file_path(
+                prefix="overcooked_visualized_state_", extension=".png"
+            )
 
         if img_path is not None:
             pygame.image.save(surface, img_path)
@@ -173,13 +274,17 @@ def render_state(self, state, grid, hud_data=None, action_probs=None):
         pygame.init()
         grid = grid or self.grid
         assert grid
-        grid_surface = pygame.surface.Surface(self._unscaled_grid_pixel_size(grid))
+        grid_surface = pygame.surface.Surface(
+            self._unscaled_grid_pixel_size(grid)
+        )
         self._render_grid(grid_surface, grid)
         self._render_players(grid_surface, state.players)
-        self._render_objects(grid_surface, state.objects, grid)       
+        self._render_objects(grid_surface, state.objects, grid)
 
         if self.scale_by_factor != 1:
-            grid_surface = scale_surface_by_factor(grid_surface, self.scale_by_factor)
+            grid_surface = scale_surface_by_factor(
+                grid_surface, self.scale_by_factor
+            )
 
         # render text after rescaling as text looks bad when is rendered small resolution and then rescalled to bigger one
         if self.is_rendering_cooking_timer:
@@ -187,22 +292,35 @@ def render_state(self, state, grid, hud_data=None, action_probs=None):
 
         # arrows does not seem good when rendered in very small resolution
         if self.is_rendering_action_probs and action_probs is not None:
-            self._render_actions_probs(grid_surface, state.players, action_probs)
+            self._render_actions_probs(
+                grid_surface, state.players, action_probs
+            )
 
         if self.is_rendering_hud and hud_data:
             hud_width = self.width or grid_surface.get_width()
-            hud_surface = pygame.surface.Surface((hud_width, self._calculate_hud_height(hud_data)))
+            hud_surface = pygame.surface.Surface(
+                (hud_width, self._calculate_hud_height(hud_data))
+            )
             hud_surface.fill(self.background_color)
             self._render_hud_data(hud_surface, hud_data)
-            rendered_surface = vstack_surfaces([hud_surface, grid_surface], self.background_color)
+            rendered_surface = vstack_surfaces(
+                [hud_surface, grid_surface], self.background_color
+            )
         else:
             hud_width = None
             rendered_surface = grid_surface
 
-        result_surface_size = (self.width or rendered_surface.get_width(), self.height or rendered_surface.get_height())
+        result_surface_size = (
+            self.width or rendered_surface.get_width(),
+            self.height or rendered_surface.get_height(),
+        )
 
         if result_surface_size != rendered_surface.get_size():
-            result_surface = blit_on_new_surface_of_size(rendered_surface, result_surface_size, background_color=self.background_color)
+            result_surface = blit_on_new_surface_of_size(
+                rendered_surface,
+                result_surface_size,
+                background_color=self.background_color,
+            )
         else:
             result_surface = rendered_surface
 
@@ -210,7 +328,7 @@ def render_state(self, state, grid, hud_data=None, action_probs=None):
 
     @property
     def scale_by_factor(self):
-        return self.tile_size/StateVisualizer.UNSCALED_TILE_SIZE
+        return self.tile_size / StateVisualizer.UNSCALED_TILE_SIZE
 
     @property
     def hud_line_height(self):
@@ -218,41 +336,53 @@ def hud_line_height(self):
 
     @staticmethod
     def _check_config_validity(config):
-        assert set(config.keys()).issubset(set(StateVisualizer.DEFAULT_VALUES.keys()))
+        assert set(config.keys()).issubset(
+            set(StateVisualizer.DEFAULT_VALUES.keys())
+        )
 
-    def _init_font(self, font_size, font_path=None, system_font_name=None):   
+    def _init_font(self, font_size, font_path=None, system_font_name=None):
         if system_font_name:
-            key = "%i-sys:%s" %(font_size, system_font_name)
-            font = self._fonts.get(key) or pygame.font.SysFont(system_font_name, font_size)
+            key = "%i-sys:%s" % (font_size, system_font_name)
+            font = self._fonts.get(key) or pygame.font.SysFont(
+                system_font_name, font_size
+            )
         else:
-            key = "%i-path:%s" %(font_size, font_path)
-            font = self._fonts.get(key) or pygame.font.Font(font_path, font_size)
+            key = "%i-path:%s" % (font_size, font_path)
+            font = self._fonts.get(key) or pygame.font.Font(
+                font_path, font_size
+            )
         self._fonts[key] = font
         return font
 
     def _unscaled_grid_pixel_size(self, grid):
         y_tiles = len(grid)
         x_tiles = len(grid[0])
-        return (x_tiles * self.UNSCALED_TILE_SIZE, y_tiles * self.UNSCALED_TILE_SIZE)
+        return (
+            x_tiles * self.UNSCALED_TILE_SIZE,
+            y_tiles * self.UNSCALED_TILE_SIZE,
+        )
 
     def _render_grid(self, surface, grid):
         for y_tile, row in enumerate(grid):
             for x_tile, tile in enumerate(row):
-                self.TERRAINS_IMG.blit_on_surface(surface, self._position_in_unscaled_pixels((x_tile, y_tile)), 
-                                             StateVisualizer.TILE_TO_FRAME_NAME[tile])
+                self.TERRAINS_IMG.blit_on_surface(
+                    surface,
+                    self._position_in_unscaled_pixels((x_tile, y_tile)),
+                    StateVisualizer.TILE_TO_FRAME_NAME[tile],
+                )
 
     def _position_in_unscaled_pixels(self, position):
         """
         get x and y coordinates in tiles, returns x and y coordinates in pixels
         """
-        (x,y) = position
+        (x, y) = position
         return (self.UNSCALED_TILE_SIZE * x, self.UNSCALED_TILE_SIZE * y)
 
     def _position_in_scaled_pixels(self, position):
         """
         get x and y coordinates in tiles, returns x and y coordinates in pixels
         """
-        (x,y) = position
+        (x, y) = position
         return (self.tile_size * x, self.tile_size * y)
 
     def _render_players(self, surface, players):
@@ -263,7 +393,7 @@ def chef_frame_name(direction_name, held_object_name):
             return frame_name
 
         def hat_frame_name(direction_name, player_color_name):
-            return "%s-%shat" %(direction_name, player_color_name)
+            return "%s-%shat" % (direction_name, player_color_name)
 
         for player_num, player in enumerate(players):
             player_color_name = self.player_colors[player_num]
@@ -281,14 +411,26 @@ def hat_frame_name(direction_name, player_color_name):
                 else:
                     held_object_name = held_obj.name
 
-            self.CHEFS_IMG.blit_on_surface(surface, self._position_in_unscaled_pixels(player.position), chef_frame_name(direction_name, held_object_name))
-            self.CHEFS_IMG.blit_on_surface(surface, self._position_in_unscaled_pixels(player.position), hat_frame_name(direction_name, player_color_name))
+            self.CHEFS_IMG.blit_on_surface(
+                surface,
+                self._position_in_unscaled_pixels(player.position),
+                chef_frame_name(direction_name, held_object_name),
+            )
+            self.CHEFS_IMG.blit_on_surface(
+                surface,
+                self._position_in_unscaled_pixels(player.position),
+                hat_frame_name(direction_name, player_color_name),
+            )
 
     @staticmethod
     def _soup_frame_name(ingredients_names, status):
-            num_onions = ingredients_names.count("onion")
-            num_tomatoes = ingredients_names.count("tomato")
-            return "soup_%s_tomato_%i_onion_%i" %(status, num_tomatoes, num_onions)
+        num_onions = ingredients_names.count("onion")
+        num_tomatoes = ingredients_names.count("tomato")
+        return "soup_%s_tomato_%i_onion_%i" % (
+            status,
+            num_tomatoes,
+            num_onions,
+        )
 
     def _render_objects(self, surface, objects, grid):
         def render_soup(surface, obj, grid):
@@ -298,28 +440,55 @@ def render_soup(surface, obj, grid):
                     soup_status = "cooked"
                 else:
                     soup_status = "idle"
-            else: # grid[x][y] != POT
+            else:  # grid[x][y] != POT
                 soup_status = "done"
-            frame_name = StateVisualizer._soup_frame_name(obj.ingredients, soup_status)
-            self.SOUPS_IMG.blit_on_surface(surface, self._position_in_unscaled_pixels(obj.position), frame_name)
+            frame_name = StateVisualizer._soup_frame_name(
+                obj.ingredients, soup_status
+            )
+            self.SOUPS_IMG.blit_on_surface(
+                surface,
+                self._position_in_unscaled_pixels(obj.position),
+                frame_name,
+            )
 
         for obj in objects.values():
             if obj.name == "soup":
                 render_soup(surface, obj, grid)
             else:
-                self.OBJECTS_IMG.blit_on_surface(surface, self._position_in_unscaled_pixels(obj.position), obj.name)            
+                self.OBJECTS_IMG.blit_on_surface(
+                    surface,
+                    self._position_in_unscaled_pixels(obj.position),
+                    obj.name,
+                )
 
     def _render_cooking_timers(self, surface, objects, grid):
         for key, obj in objects.items():
             (x_pos, y_pos) = obj.position
             if obj.name == "soup" and grid[y_pos][x_pos] == POT:
-                if obj._cooking_tick != -1 and (obj._cooking_tick <= obj.cook_time or self.show_timer_when_cooked):
-                    text_surface = self.cooking_timer_font.render(str(obj._cooking_tick), True, self.cooking_timer_font_color)
-                    (tile_pos_x, tile_pos_y) = self._position_in_scaled_pixels(obj.position)
+                if obj._cooking_tick != -1 and (
+                    obj._cooking_tick <= obj.cook_time
+                    or self.show_timer_when_cooked
+                ):
+                    text_surface = self.cooking_timer_font.render(
+                        str(obj._cooking_tick),
+                        True,
+                        self.cooking_timer_font_color,
+                    )
+                    (tile_pos_x, tile_pos_y) = self._position_in_scaled_pixels(
+                        obj.position
+                    )
 
                     # calculate font position to be in center on x axis, and 0.9 from top on y axis
-                    font_position = (tile_pos_x + int((self.tile_size - text_surface.get_width() ) * 0.5), 
-                                 tile_pos_y + int((self.tile_size - text_surface.get_height() ) * 0.9))
+                    font_position = (
+                        tile_pos_x
+                        + int(
+                            (self.tile_size - text_surface.get_width()) * 0.5
+                        ),
+                        tile_pos_y
+                        + int(
+                            (self.tile_size - text_surface.get_height()) * 0.9
+                        ),
+                    )
                     surface.blit(text_surface, font_position)
 
     def _sorted_hud_items(self, hud_data):
@@ -330,6 +499,7 @@ def default_order_then_alphabetic(item):
             except:
                 i = 99999
             return (i, key)
+
         return sorted(hud_data.items(), key=default_order_then_alphabetic)
 
     def _key_to_hud_text(self, key):
@@ -337,7 +507,10 @@ def _key_to_hud_text(self, key):
 
     def _render_hud_data(self, surface, hud_data):
         def hud_text_position(line_num):
-            return (self.hud_margin_left, self.hud_margin_top + self.hud_line_height * line_num)
+            return (
+                self.hud_margin_left,
+                self.hud_margin_top + self.hud_line_height * line_num,
+            )
 
         def hud_recipes_position(text_surface, text_surface_position):
             (text_surface_x, text_surface_y) = text_surface_position
@@ -347,90 +520,196 @@ def get_hud_recipes_surface(orders_dicts):
             order_width = order_height = self.hud_order_size
             scaled_order_size = (order_width, order_width)
             orders_surface_height = order_height
-            orders_surface_width = len(orders_dicts) * order_width + (len(orders_dicts) - 1) * self.hud_distance_between_orders
-            unscaled_order_size = (self.UNSCALED_TILE_SIZE, self.UNSCALED_TILE_SIZE)
-
-            recipes_surface = pygame.surface.Surface((orders_surface_width, orders_surface_height))
+            orders_surface_width = (
+                len(orders_dicts) * order_width
+                + (len(orders_dicts) - 1) * self.hud_distance_between_orders
+            )
+            unscaled_order_size = (
+                self.UNSCALED_TILE_SIZE,
+                self.UNSCALED_TILE_SIZE,
+            )
+
+            recipes_surface = pygame.surface.Surface(
+                (orders_surface_width, orders_surface_height)
+            )
             recipes_surface.fill(self.background_color)
             next_surface_x = 0
             for order_dict in orders_dicts:
-                frame_name = StateVisualizer._soup_frame_name(order_dict["ingredients"], "done")
-                unscaled_order_surface = pygame.surface.Surface(unscaled_order_size)
+                frame_name = StateVisualizer._soup_frame_name(
+                    order_dict["ingredients"], "done"
+                )
+                unscaled_order_surface = pygame.surface.Surface(
+                    unscaled_order_size
+                )
                 unscaled_order_surface.fill(self.background_color)
-                self.SOUPS_IMG.blit_on_surface(unscaled_order_surface, (0,0), frame_name)
+                self.SOUPS_IMG.blit_on_surface(
+                    unscaled_order_surface, (0, 0), frame_name
+                )
                 if scaled_order_size == unscaled_order_size:
-                    scaled_order_surface = unscaled_order_surface 
+                    scaled_order_surface = unscaled_order_surface
                 else:
-                    scaled_order_surface = pygame.transform.scale(unscaled_order_surface, (order_width, order_width))
+                    scaled_order_surface = pygame.transform.scale(
+                        unscaled_order_surface, (order_width, order_width)
+                    )
                 recipes_surface.blit(scaled_order_surface, (next_surface_x, 0))
-                next_surface_x += order_width + self.hud_distance_between_orders
+                next_surface_x += (
+                    order_width + self.hud_distance_between_orders
+                )
             return recipes_surface
 
-        for hud_line_num, (key, value) in enumerate(self._sorted_hud_items(hud_data)):
+        for hud_line_num, (key, value) in enumerate(
+            self._sorted_hud_items(hud_data)
+        ):
             hud_text = self._key_to_hud_text(key)
-            if key not in ["all_orders", "bonus_orders", "start_all_orders", "start_bonus_orders"]:
+            if key not in [
+                "all_orders",
+                "bonus_orders",
+                "start_all_orders",
+                "start_bonus_orders",
+            ]:
                 hud_text += str(value)
 
-            text_surface = self.hud_font.render(hud_text, True, self.hud_font_color)
+            text_surface = self.hud_font.render(
+                hud_text, True, self.hud_font_color
+            )
             text_surface_position = hud_text_position(hud_line_num)
             surface.blit(text_surface, text_surface_position)
 
-            if key in ["all_orders", "bonus_orders", "start_all_orders", "start_bonus_orders"] and value:
-                recipes_surface_position = hud_recipes_position(text_surface, text_surface_position)
+            if (
+                key
+                in [
+                    "all_orders",
+                    "bonus_orders",
+                    "start_all_orders",
+                    "start_bonus_orders",
+                ]
+                and value
+            ):
+                recipes_surface_position = hud_recipes_position(
+                    text_surface, text_surface_position
+                )
                 recipes_surface = get_hud_recipes_surface(value)
-                assert recipes_surface.get_width() + text_surface.get_width() <= surface.get_width(), "surface width is too small to fit recipes in single line"
+                assert (
+                    recipes_surface.get_width() + text_surface.get_width()
+                    <= surface.get_width()
+                ), "surface width is too small to fit recipes in single line"
                 surface.blit(recipes_surface, recipes_surface_position)
 
     def _calculate_hud_height(self, hud_data):
-        return self.hud_margin_top + len(hud_data) * self.hud_line_height + self.hud_margin_bottom
-
-    def _render_on_tile_position(self, scaled_grid_surface, source_surface, tile_position, horizontal_align="left", vertical_align="top"):
+        return (
+            self.hud_margin_top
+            + len(hud_data) * self.hud_line_height
+            + self.hud_margin_bottom
+        )
+
+    def _render_on_tile_position(
+        self,
+        scaled_grid_surface,
+        source_surface,
+        tile_position,
+        horizontal_align="left",
+        vertical_align="top",
+    ):
         assert vertical_align in ["top", "center", "bottom"]
         left_x, top_y = self._position_in_scaled_pixels(tile_position)
         if horizontal_align == "left":
             x = left_x
         elif horizontal_align == "center":
-            x = left_x + (self.tile_size - source_surface.get_width())/2
+            x = left_x + (self.tile_size - source_surface.get_width()) / 2
         elif horizontal_align == "right":
             x = left_x + self.tile_size - source_surface.get_width()
         else:
-            raise ValueError("horizontal_align can have one of the values: "+str(["left", "center", "right"]))
+            raise ValueError(
+                "horizontal_align can have one of the values: "
+                + str(["left", "center", "right"])
+            )
 
         if vertical_align == "top":
             y = top_y
         elif vertical_align == "center":
-            y = top_y + (self.tile_size - source_surface.get_height())/2
+            y = top_y + (self.tile_size - source_surface.get_height()) / 2
         elif vertical_align == "bottom":
             y = top_y + self.tile_size - source_surface.get_height()
         else:
-            raise ValueError("vertical_align can have one of the values: "+str(["top", "center", "bottom"]))
+            raise ValueError(
+                "vertical_align can have one of the values: "
+                + str(["top", "center", "bottom"])
+            )
 
         scaled_grid_surface.blit(source_surface, (x, y))
 
     def _render_actions_probs(self, surface, players, action_probs):
-        direction_to_rotation = {Direction.NORTH:0, Direction.WEST:90 , Direction.SOUTH:180, Direction.EAST:270}
+        direction_to_rotation = {
+            Direction.NORTH: 0,
+            Direction.WEST: 90,
+            Direction.SOUTH: 180,
+            Direction.EAST: 270,
+        }
         direction_to_aligns = {
-            Direction.NORTH: {"horizontal_align": "center", "vertical_align":"bottom"},
-            Direction.WEST: {"horizontal_align": "right", "vertical_align":"center"},
-            Direction.SOUTH: {"horizontal_align": "center", "vertical_align":"top"},
-            Direction.EAST: {"horizontal_align": "left", "vertical_align":"center"}}
-
-        rescaled_arrow = pygame.transform.scale(self.ARROW_IMG, (self.tile_size, self.tile_size))
+            Direction.NORTH: {
+                "horizontal_align": "center",
+                "vertical_align": "bottom",
+            },
+            Direction.WEST: {
+                "horizontal_align": "right",
+                "vertical_align": "center",
+            },
+            Direction.SOUTH: {
+                "horizontal_align": "center",
+                "vertical_align": "top",
+            },
+            Direction.EAST: {
+                "horizontal_align": "left",
+                "vertical_align": "center",
+            },
+        }
+
+        rescaled_arrow = pygame.transform.scale(
+            self.ARROW_IMG, (self.tile_size, self.tile_size)
+        )
         # divide width by math.sqrt(2) to always fit both interact icon and stay icon into single tile
-        rescaled_interact = pygame.transform.scale(self.INTERACT_IMG, (int(self.tile_size/math.sqrt(2)), self.tile_size))
-        rescaled_stay = pygame.transform.scale(self.STAY_IMG, (int(self.tile_size/math.sqrt(2)), self.tile_size))
+        rescaled_interact = pygame.transform.scale(
+            self.INTERACT_IMG,
+            (int(self.tile_size / math.sqrt(2)), self.tile_size),
+        )
+        rescaled_stay = pygame.transform.scale(
+            self.STAY_IMG, (int(self.tile_size / math.sqrt(2)), self.tile_size)
+        )
         for player, probs in zip(players, action_probs):
             if probs is not None:
                 for action in Action.ALL_ACTIONS:
                     # use math sqrt to make probability proportional to area of the image
                     size = math.sqrt(probs[Action.ACTION_TO_INDEX[action]])
                     if action == "interact":
-                        img = pygame.transform.rotozoom(rescaled_interact, 0, size)
-                        self._render_on_tile_position(surface, img, player.position, horizontal_align="left", vertical_align="center")
+                        img = pygame.transform.rotozoom(
+                            rescaled_interact, 0, size
+                        )
+                        self._render_on_tile_position(
+                            surface,
+                            img,
+                            player.position,
+                            horizontal_align="left",
+                            vertical_align="center",
+                        )
                     elif action == Action.STAY:
                         img = pygame.transform.rotozoom(rescaled_stay, 0, size)
-                        self._render_on_tile_position(surface, img, player.position, horizontal_align="right", vertical_align="center")
+                        self._render_on_tile_position(
+                            surface,
+                            img,
+                            player.position,
+                            horizontal_align="right",
+                            vertical_align="center",
+                        )
                     else:
-                        position = Action.move_in_direction(player.position, action)
-                        img =  pygame.transform.rotozoom(rescaled_arrow, direction_to_rotation[action], size)
-                        self._render_on_tile_position(surface, img, position, **direction_to_aligns[action])
\ No newline at end of file
+                        position = Action.move_in_direction(
+                            player.position, action
+                        )
+                        img = pygame.transform.rotozoom(
+                            rescaled_arrow, direction_to_rotation[action], size
+                        )
+                        self._render_on_tile_position(
+                            surface,
+                            img,
+                            position,
+                            **direction_to_aligns[action]
+                        )
diff --git a/src/overcooked_ai_py/visualization/visualization_utils.py b/src/overcooked_ai_py/visualization/visualization_utils.py
index b8974b4a..4e37e0cc 100644
--- a/src/overcooked_ai_py/visualization/visualization_utils.py
+++ b/src/overcooked_ai_py/visualization/visualization_utils.py
@@ -1,15 +1,38 @@
 from IPython.display import Image, display
-from ipywidgets import interactive, IntSlider
+from ipywidgets import IntSlider, interactive
+
 
 def show_image_in_ipython(data, *args, **kwargs):
     display(Image(data, *args, **kwargs))
 
+
 def ipython_images_slider(image_pathes_list, slider_label="", first_arg=0):
     def display_f(**kwargs):
         display(Image(image_pathes_list[kwargs[slider_label]]))
-    return interactive(display_f, **{slider_label: IntSlider(min=0, max=len(image_pathes_list)-1,step=1)})
 
-def show_ipython_images_slider(image_pathes_list, slider_label="", first_arg=0):
+    return interactive(
+        display_f,
+        **{
+            slider_label: IntSlider(
+                min=0, max=len(image_pathes_list) - 1, step=1
+            )
+        }
+    )
+
+
+def show_ipython_images_slider(
+    image_pathes_list, slider_label="", first_arg=0
+):
     def display_f(**kwargs):
         display(Image(image_pathes_list[kwargs[slider_label]]))
-    display(interactive(display_f, **{slider_label: IntSlider(min=0, max=len(image_pathes_list)-1,step=1)}))
+
+    display(
+        interactive(
+            display_f,
+            **{
+                slider_label: IntSlider(
+                    min=0, max=len(image_pathes_list) - 1, step=1
+                )
+            }
+        )
+    )


Test patch:
diff --git a/src/overcooked_ai_py/data/testing/test_state_visualizer/Create_state_visualizer_tests.ipynb b/src/overcooked_ai_py/data/testing/test_state_visualizer/Create_state_visualizer_tests.ipynb
index 24de336e..2834725a 100644
--- a/src/overcooked_ai_py/data/testing/test_state_visualizer/Create_state_visualizer_tests.ipynb
+++ b/src/overcooked_ai_py/data/testing/test_state_visualizer/Create_state_visualizer_tests.ipynb
@@ -18,10 +18,21 @@
     "# This notebook is for tests generation for state visualizer\n",
     "# if you cannot run this notebook in conda env check this out: https://medium.com/@nrk25693/how-to-add-your-conda-environment-to-your-jupyter-notebook-in-just-4-steps-abeab8b8d084\n",
     "from overcooked_ai_py.agents.benchmarking import AgentEvaluator\n",
-    "from overcooked_ai_py.agents.agent import SampleAgent, GreedyHumanModel, RandomAgent\n",
-    "from overcooked_ai_py.mdp.overcooked_mdp import OvercookedGridworld, OvercookedState, Recipe, SoupState\n",
+    "from overcooked_ai_py.agents.agent import (\n",
+    "    SampleAgent,\n",
+    "    GreedyHumanModel,\n",
+    "    RandomAgent,\n",
+    ")\n",
+    "from overcooked_ai_py.mdp.overcooked_mdp import (\n",
+    "    OvercookedGridworld,\n",
+    "    OvercookedState,\n",
+    "    Recipe,\n",
+    "    SoupState,\n",
+    ")\n",
     "from overcooked_ai_py.visualization.state_visualizer import StateVisualizer\n",
-    "from overcooked_ai_py.visualization.visualization_utils import show_image_in_ipython\n",
+    "from overcooked_ai_py.visualization.visualization_utils import (\n",
+    "    show_image_in_ipython,\n",
+    ")\n",
     "from overcooked_ai_py.utils import generate_temporary_file_path\n",
     "from overcooked_ai_py.static import FONTS_DIR\n",
     "from overcooked_ai_py.mdp.layout_generator import POT\n",
@@ -43,50 +54,66 @@
     "\n",
     "\n",
     "Recipe.configure({})\n",
+    "\n",
+    "\n",
     "def display_and_export_to_array(test_dict):\n",
     "    test_dict = copy.deepcopy(test_dict)\n",
-    "    test_dict[\"kwargs\"][\"state\"] = OvercookedState.from_dict(test_dict[\"kwargs\"][\"state\"])\n",
-    "    surface =  StateVisualizer(**test_dict[\"config\"]).render_state(**test_dict[\"kwargs\"])\n",
-    "    img_path = generate_temporary_file_path(\"temporary_visualization\", extension=\".png\")\n",
+    "    test_dict[\"kwargs\"][\"state\"] = OvercookedState.from_dict(\n",
+    "        test_dict[\"kwargs\"][\"state\"]\n",
+    "    )\n",
+    "    surface = StateVisualizer(**test_dict[\"config\"]).render_state(\n",
+    "        **test_dict[\"kwargs\"]\n",
+    "    )\n",
+    "    img_path = generate_temporary_file_path(\n",
+    "        \"temporary_visualization\", extension=\".png\"\n",
+    "    )\n",
     "    pygame.image.save(surface, img_path)\n",
     "    print(\"check if image is okay\")\n",
     "    show_image_in_ipython(img_path)\n",
     "    return pygame.surfarray.array3d(surface)\n",
     "\n",
+    "\n",
     "def save_test_files(test_dicts_list, test_result_arrays, json_filename):\n",
     "    for d, a in zip(test_dicts_list, test_result_arrays):\n",
-    "        with open(d[\"result_array_filename\"], 'wb') as f:\n",
+    "        with open(d[\"result_array_filename\"], \"wb\") as f:\n",
     "            np.save(f, a)\n",
     "    with open(json_filename, \"w\") as f:\n",
     "        json.dump(test_dicts_list, f)\n",
-    "        \n",
+    "\n",
+    "\n",
     "DEFAULT_VALUES = {\n",
-    "        \"height\": None, # if None use grid_width - NOTE: can chop down hud if hud is wider than grid\n",
-    "        \"width\": None, # if None use (hud_height+grid_height)\n",
-    "        \"tile_size\": 75,\n",
-    "        \"window_fps\": 30,\n",
-    "        \"player_colors\": ['blue', 'green'],\n",
-    "        \"is_rendering_hud\": True,\n",
-    "        \"hud_font_size\": 10,\n",
-    "        \"hud_system_font_name\": None, # if set to None use hud_font_path\n",
-    "        # needs to be overwritten with default - every pc has different pathes \"hud_font_path\": roboto_path, \n",
-    "        \"hud_font_color\": (255,255,255), # white\n",
-    "        \"hud_data_default_key_order\" : [\"all_orders\", \"bonus_orders\", \"time_left\", \"score\", \"potential\"],\n",
-    "        \"hud_interline_size\": 10,\n",
-    "        \"hud_margin_bottom\": 10,\n",
-    "        \"hud_margin_top\": 10,\n",
-    "        \"hud_margin_left\": 10,\n",
-    "        \"hud_distance_between_orders\": 5,\n",
-    "        \"hud_order_size\": 15,\n",
-    "        \"is_rendering_cooking_timer\": True,\n",
-    "        \"show_timer_when_cooked\": True,\n",
-    "        \"cooking_timer_font_size\": 20, # # if set to None use cooking_timer_font_path\n",
-    "        # needs to be overwritten with default - every pc has different pathes \"cooking_timer_font_path\": roboto_path, \n",
-    "        \"cooking_timer_system_font_name\": None,\n",
-    "        \"cooking_timer_font_color\": (255, 0, 0), # red\n",
-    "        \"grid\": None,\n",
-    "        \"background_color\": (155, 101, 0) # color of empty counter\n",
-    "    }\n"
+    "    \"height\": None,  # if None use grid_width - NOTE: can chop down hud if hud is wider than grid\n",
+    "    \"width\": None,  # if None use (hud_height+grid_height)\n",
+    "    \"tile_size\": 75,\n",
+    "    \"window_fps\": 30,\n",
+    "    \"player_colors\": [\"blue\", \"green\"],\n",
+    "    \"is_rendering_hud\": True,\n",
+    "    \"hud_font_size\": 10,\n",
+    "    \"hud_system_font_name\": None,  # if set to None use hud_font_path\n",
+    "    # needs to be overwritten with default - every pc has different pathes \"hud_font_path\": roboto_path,\n",
+    "    \"hud_font_color\": (255, 255, 255),  # white\n",
+    "    \"hud_data_default_key_order\": [\n",
+    "        \"all_orders\",\n",
+    "        \"bonus_orders\",\n",
+    "        \"time_left\",\n",
+    "        \"score\",\n",
+    "        \"potential\",\n",
+    "    ],\n",
+    "    \"hud_interline_size\": 10,\n",
+    "    \"hud_margin_bottom\": 10,\n",
+    "    \"hud_margin_top\": 10,\n",
+    "    \"hud_margin_left\": 10,\n",
+    "    \"hud_distance_between_orders\": 5,\n",
+    "    \"hud_order_size\": 15,\n",
+    "    \"is_rendering_cooking_timer\": True,\n",
+    "    \"show_timer_when_cooked\": True,\n",
+    "    \"cooking_timer_font_size\": 20,  # # if set to None use cooking_timer_font_path\n",
+    "    # needs to be overwritten with default - every pc has different pathes \"cooking_timer_font_path\": roboto_path,\n",
+    "    \"cooking_timer_system_font_name\": None,\n",
+    "    \"cooking_timer_font_color\": (255, 0, 0),  # red\n",
+    "    \"grid\": None,\n",
+    "    \"background_color\": (155, 101, 0),  # color of empty counter\n",
+    "}"
    ]
   },
   {
@@ -172,7 +199,9 @@
     "mdp = OvercookedGridworld.from_layout_name(layout_name=\"cramped_room\")\n",
     "agent_eval = AgentEvaluator(env_params={\"horizon\": 1}, mdp_fn=lambda _: mdp)\n",
     "\n",
-    "trajectory_random_pair = agent_eval.evaluate_random_pair(num_games=1, display=False, native_eval=True)\n",
+    "trajectory_random_pair = agent_eval.evaluate_random_pair(\n",
+    "    num_games=1, display=False, native_eval=True\n",
+    ")\n",
     "grid = trajectory_random_pair[\"mdp_params\"][0][\"terrain\"]\n",
     "state = trajectory_random_pair[\"ep_states\"][0][0]\n",
     "\n",
@@ -180,11 +209,14 @@
     "config[\"tile_size\"] = 30\n",
     "config[\"hud_font_size\"] = 20\n",
     "hud_data = {\"time_left\": 100}\n",
-    "kwargs = {\"hud_data\": hud_data, \"grid\":grid, \"state\":state.to_dict()}\n",
+    "kwargs = {\"hud_data\": hud_data, \"grid\": grid, \"state\": state.to_dict()}\n",
     "\n",
-    "test_hud_1 = {\"config\": config, \"kwargs\": kwargs, \n",
-    "             \"comment\": \"Test simple text in hud. NOTE: failing to render stuff outside HUD also fails this test\",\n",
-    "             \"result_array_filename\": \"test_hud_1.npy\"}\n",
+    "test_hud_1 = {\n",
+    "    \"config\": config,\n",
+    "    \"kwargs\": kwargs,\n",
+    "    \"comment\": \"Test simple text in hud. NOTE: failing to render stuff outside HUD also fails this test\",\n",
+    "    \"result_array_filename\": \"test_hud_1.npy\",\n",
+    "}\n",
     "print(\"testing test_hud_1 dict\")\n",
     "test_hud_1_array = display_and_export_to_array(test_hud_1)\n",
     "hud_tests.append(test_hud_1)\n",
@@ -195,11 +227,16 @@
     "config[\"tile_size\"] = 30\n",
     "config[\"hud_font_size\"] = 20\n",
     "config[\"hud_order_size\"] = 25\n",
-    "hud_data = {\"all_orders\": trajectory_random_pair[\"mdp_params\"][0][\"start_all_orders\"]}\n",
-    "kwargs = {\"hud_data\": hud_data, \"grid\":grid, \"state\":state.to_dict()}\n",
-    "test_hud_2 = {\"config\": config, \"kwargs\": kwargs, \n",
-    "             \"comment\": \"Test simple recipes in hud. NOTE: failing to render stuff outside HUD also fails this test\",\n",
-    "             \"result_array_filename\": \"test_hud_2.npy\"}\n",
+    "hud_data = {\n",
+    "    \"all_orders\": trajectory_random_pair[\"mdp_params\"][0][\"start_all_orders\"]\n",
+    "}\n",
+    "kwargs = {\"hud_data\": hud_data, \"grid\": grid, \"state\": state.to_dict()}\n",
+    "test_hud_2 = {\n",
+    "    \"config\": config,\n",
+    "    \"kwargs\": kwargs,\n",
+    "    \"comment\": \"Test simple recipes in hud. NOTE: failing to render stuff outside HUD also fails this test\",\n",
+    "    \"result_array_filename\": \"test_hud_2.npy\",\n",
+    "}\n",
     "\n",
     "print(\"testing test_hud_2 dict\")\n",
     "test_hud_2_array = display_and_export_to_array(test_hud_2)\n",
@@ -209,11 +246,14 @@
     "config = copy.deepcopy(DEFAULT_VALUES)\n",
     "config[\"tile_size\"] = 30\n",
     "config[\"hud_font_size\"] = 20\n",
-    "hud_data = {\"time_left\": 100, \"a\":\"a\", \"b\": \"b\", \"c\": \"c\",  \"z\":\"z\"}\n",
-    "kwargs = {\"hud_data\": hud_data, \"grid\":grid, \"state\":state.to_dict()}\n",
-    "test_hud_3 = {\"config\": config, \"kwargs\": kwargs, \n",
-    "             \"comment\": \"Test chronology of hud based on keys NOTE: failing to render stuff outside HUD also fails this test\",\n",
-    "             \"result_array_filename\": \"test_hud_3.npy\"}\n",
+    "hud_data = {\"time_left\": 100, \"a\": \"a\", \"b\": \"b\", \"c\": \"c\", \"z\": \"z\"}\n",
+    "kwargs = {\"hud_data\": hud_data, \"grid\": grid, \"state\": state.to_dict()}\n",
+    "test_hud_3 = {\n",
+    "    \"config\": config,\n",
+    "    \"kwargs\": kwargs,\n",
+    "    \"comment\": \"Test chronology of hud based on keys NOTE: failing to render stuff outside HUD also fails this test\",\n",
+    "    \"result_array_filename\": \"test_hud_3.npy\",\n",
+    "}\n",
     "print(\"testing test_hud_3 dict\")\n",
     "test_hud_3_array = display_and_export_to_array(test_hud_3)\n",
     "hud_tests.append(test_hud_3)\n",
@@ -278,10 +318,13 @@
     "size_test_results = []\n",
     "config = copy.deepcopy(DEFAULT_VALUES)\n",
     "config[\"tile_size\"] = 60\n",
-    "kwargs = {\"hud_data\": {}, \"grid\":grid, \"state\":state.to_dict()}\n",
-    "test_size_1 = {\"config\": config, \"kwargs\": kwargs, \n",
-    "             \"comment\": \"Test of setting of tile_size (no hud)\",\n",
-    "             \"result_array_filename\": \"test_size_1.npy\"}\n",
+    "kwargs = {\"hud_data\": {}, \"grid\": grid, \"state\": state.to_dict()}\n",
+    "test_size_1 = {\n",
+    "    \"config\": config,\n",
+    "    \"kwargs\": kwargs,\n",
+    "    \"comment\": \"Test of setting of tile_size (no hud)\",\n",
+    "    \"result_array_filename\": \"test_size_1.npy\",\n",
+    "}\n",
     "print(\"testing test_size_1 dict\")\n",
     "test_size_1_array = display_and_export_to_array(test_size_1)\n",
     "size_tests.append(test_size_1)\n",
@@ -291,10 +334,13 @@
     "config[\"tile_size\"] = 15\n",
     "config[\"height\"] = 100\n",
     "config[\"width\"] = 100\n",
-    "kwargs = {\"hud_data\": {}, \"grid\":grid, \"state\":state.to_dict()}\n",
-    "test_size_2 = {\"config\": config, \"kwargs\": kwargs, \n",
-    "             \"comment\": \"Test of setting of width and height (no hud)\",\n",
-    "             \"result_array_filename\": \"test_size_2.npy\"}\n",
+    "kwargs = {\"hud_data\": {}, \"grid\": grid, \"state\": state.to_dict()}\n",
+    "test_size_2 = {\n",
+    "    \"config\": config,\n",
+    "    \"kwargs\": kwargs,\n",
+    "    \"comment\": \"Test of setting of width and height (no hud)\",\n",
+    "    \"result_array_filename\": \"test_size_2.npy\",\n",
+    "}\n",
     "print(\"testing test_size_2 dict\")\n",
     "test_size_2_array = display_and_export_to_array(test_size_2)\n",
     "size_tests.append(test_size_2)\n",
@@ -308,7 +354,9 @@
    "outputs": [],
    "source": [
     "# if all stuff above is ok run this cell\n",
-    "save_test_files(size_tests, size_test_results, \"render_state_data_test_sizes.json\")"
+    "save_test_files(\n",
+    "    size_tests, size_test_results, \"render_state_data_test_sizes.json\"\n",
+    ")"
    ]
   },
   {
@@ -358,17 +406,26 @@
     "\n",
     "while not rendered_cooking_timer:\n",
     "    mdp = OvercookedGridworld.from_layout_name(layout_name=\"cramped_room\")\n",
-    "    agent_eval = AgentEvaluator(env_params={\"horizon\": 1000}, mdp_fn=lambda _: mdp)\n",
-    "    trajectory_random_pair = agent_eval.evaluate_random_pair(num_games=1, display=False, native_eval=True)\n",
+    "    agent_eval = AgentEvaluator(\n",
+    "        env_params={\"horizon\": 1000}, mdp_fn=lambda _: mdp\n",
+    "    )\n",
+    "    trajectory_random_pair = agent_eval.evaluate_random_pair(\n",
+    "        num_games=1, display=False, native_eval=True\n",
+    "    )\n",
     "    grid = trajectory_random_pair[\"mdp_params\"][0][\"terrain\"]\n",
     "    for state in trajectory_random_pair[\"ep_states\"][0]:\n",
     "        if not rendered_cooking_timer and has_cooking_timer(state, grid):\n",
-    "            kwargs = {\"hud_data\": {}, \"grid\":grid, \"state\":state.to_dict()}\n",
-    "            test_cooking_timer_1 = {\"config\": config, \"kwargs\": kwargs, \n",
-    "             \"comment\": \"Test of setting rendering cooking timer\",\n",
-    "             \"result_array_filename\": \"test_cooking_timer_display_1.npy\"}\n",
+    "            kwargs = {\"hud_data\": {}, \"grid\": grid, \"state\": state.to_dict()}\n",
+    "            test_cooking_timer_1 = {\n",
+    "                \"config\": config,\n",
+    "                \"kwargs\": kwargs,\n",
+    "                \"comment\": \"Test of setting rendering cooking timer\",\n",
+    "                \"result_array_filename\": \"test_cooking_timer_display_1.npy\",\n",
+    "            }\n",
     "            print(\"test_cooking_timer_1 dict\")\n",
-    "            test_test_cooking_timer_1_array = display_and_export_to_array(test_cooking_timer_1)\n",
+    "            test_test_cooking_timer_1_array = display_and_export_to_array(\n",
+    "                test_cooking_timer_1\n",
+    "            )\n",
     "            rendered_cooking_timer = True"
    ]
   },
@@ -379,7 +436,11 @@
    "outputs": [],
    "source": [
     "# if all stuff above is ok run this cell\n",
-    "save_test_files([test_cooking_timer_1], [test_test_cooking_timer_1_array], \"render_state_data_test_cooking_display.json\")"
+    "save_test_files(\n",
+    "    [test_cooking_timer_1],\n",
+    "    [test_test_cooking_timer_1_array],\n",
+    "    \"render_state_data_test_cooking_display.json\",\n",
+    ")"
    ]
   },
   {
@@ -610,19 +671,23 @@
     "config[\"cooking_timer_font_size\"] = 15\n",
     "mdp = OvercookedGridworld.from_layout_name(layout_name=\"cramped_room_tomato\")\n",
     "agent_eval = AgentEvaluator(env_params={\"horizon\": 1001}, mdp_fn=lambda _: mdp)\n",
-    "trajectory_random_pair = agent_eval.evaluate_random_pair(num_games=1, display=False, native_eval=True)\n",
+    "trajectory_random_pair = agent_eval.evaluate_random_pair(\n",
+    "    num_games=1, display=False, native_eval=True\n",
+    ")\n",
     "grid = trajectory_random_pair[\"mdp_params\"][0][\"terrain\"]\n",
     "for i in range(1, 11):\n",
-    "    state = trajectory_random_pair[\"ep_states\"][0][i*100]\n",
-    "    kwargs = {\"hud_data\": {}, \"grid\":grid, \"state\":state.to_dict()}\n",
-    "    test_dict = {\"config\": config, \"kwargs\": kwargs, \n",
-    "     \"comment\": \"Various tests\",\n",
-    "     \"result_array_filename\": \"test_various_display_%d.npy\"%i}\n",
-    "    print(\"test_various_display_%i dict\"%i)\n",
+    "    state = trajectory_random_pair[\"ep_states\"][0][i * 100]\n",
+    "    kwargs = {\"hud_data\": {}, \"grid\": grid, \"state\": state.to_dict()}\n",
+    "    test_dict = {\n",
+    "        \"config\": config,\n",
+    "        \"kwargs\": kwargs,\n",
+    "        \"comment\": \"Various tests\",\n",
+    "        \"result_array_filename\": \"test_various_display_%d.npy\" % i,\n",
+    "    }\n",
+    "    print(\"test_various_display_%i dict\" % i)\n",
     "    test_array = display_and_export_to_array(test_dict)\n",
     "    various_tests.append(test_dict)\n",
-    "    various_results.append(test_array)\n",
-    "    "
+    "    various_results.append(test_array)"
    ]
   },
   {
@@ -632,7 +697,9 @@
    "outputs": [],
    "source": [
     "# if all stuff above is ok run this cell\n",
-    "save_test_files(various_tests, various_results, \"render_state_data_test_various.json\")"
+    "save_test_files(\n",
+    "    various_tests, various_results, \"render_state_data_test_various.json\"\n",
+    ")"
    ]
   },
   {
@@ -855,6 +922,7 @@
    ],
    "source": [
     "from overcooked_ai_py.mdp.layout_generator import LayoutGenerator\n",
+    "\n",
     "# testing some states hoping it can find unexpected bugs\n",
     "generated_layout_tests = []\n",
     "generated_layout_results = []\n",
@@ -867,26 +935,31 @@
     "    \"inner_shape\": (7, 5),\n",
     "    \"prop_empty\": 0.7,\n",
     "    \"prop_feats\": 0.4,\n",
-    "    \"start_all_orders\" : [\n",
-    "        { \"ingredients\" : [\"onion\", \"onion\", \"onion\"]}\n",
-    "    ],\n",
-    "    \"recipe_values\" : [20],\n",
-    "    \"recipe_times\" : [20],\n",
-    "    \"display\": False\n",
+    "    \"start_all_orders\": [{\"ingredients\": [\"onion\", \"onion\", \"onion\"]}],\n",
+    "    \"recipe_values\": [20],\n",
+    "    \"recipe_times\": [20],\n",
+    "    \"display\": False,\n",
     "}\n",
     "\n",
-    "mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(DEFAULT_MDP_GEN_PARAMS, outer_shape=(7, 5))\n",
+    "mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(\n",
+    "    DEFAULT_MDP_GEN_PARAMS, outer_shape=(7, 5)\n",
+    ")\n",
     "\n",
     "agent_eval = AgentEvaluator({\"horizon\": 1001}, mdp_fn)\n",
-    "trajectory_random_pair = agent_eval.evaluate_random_pair(num_games=1, display=False, native_eval=True)\n",
+    "trajectory_random_pair = agent_eval.evaluate_random_pair(\n",
+    "    num_games=1, display=False, native_eval=True\n",
+    ")\n",
     "grid = trajectory_random_pair[\"mdp_params\"][0][\"terrain\"]\n",
     "for i in range(1, 11):\n",
-    "    state = trajectory_random_pair[\"ep_states\"][0][i*100]\n",
-    "    kwargs = {\"hud_data\": {}, \"grid\":grid, \"state\":state.to_dict()}\n",
-    "    test_dict = {\"config\": config, \"kwargs\": kwargs, \n",
-    "     \"comment\": \"Various tests on generated layout\",\n",
-    "     \"result_array_filename\": \"test_generated_layout_display_%d.npy\"%i}\n",
-    "    print(\"test_generated_layout_display_%i dict\"%i)\n",
+    "    state = trajectory_random_pair[\"ep_states\"][0][i * 100]\n",
+    "    kwargs = {\"hud_data\": {}, \"grid\": grid, \"state\": state.to_dict()}\n",
+    "    test_dict = {\n",
+    "        \"config\": config,\n",
+    "        \"kwargs\": kwargs,\n",
+    "        \"comment\": \"Various tests on generated layout\",\n",
+    "        \"result_array_filename\": \"test_generated_layout_display_%d.npy\" % i,\n",
+    "    }\n",
+    "    print(\"test_generated_layout_display_%i dict\" % i)\n",
     "    test_array = display_and_export_to_array(test_dict)\n",
     "    generated_layout_tests.append(test_dict)\n",
     "    generated_layout_results.append(test_array)"
@@ -899,7 +972,11 @@
    "outputs": [],
    "source": [
     "# if all stuff above is ok run this cell\n",
-    "save_test_files(generated_layout_tests, generated_layout_results, \"render_state_data_test_generated_layout.json\")"
+    "save_test_files(\n",
+    "    generated_layout_tests,\n",
+    "    generated_layout_results,\n",
+    "    \"render_state_data_test_generated_layout.json\",\n",
+    ")"
    ]
   },
   {
@@ -1113,26 +1190,47 @@
     "config[\"cooking_timer_font_size\"] = 15\n",
     "mdp = OvercookedGridworld.from_layout_name(layout_name=\"cramped_room\")\n",
     "agent_eval = AgentEvaluator(env_params={\"horizon\": 1001}, mdp_fn=lambda _: mdp)\n",
-    "trajectory_random_pair = agent_eval.evaluate_random_pair(num_games=1, display=False, native_eval=True)\n",
+    "trajectory_random_pair = agent_eval.evaluate_random_pair(\n",
+    "    num_games=1, display=False, native_eval=True\n",
+    ")\n",
     "grid = trajectory_random_pair[\"mdp_params\"][0][\"terrain\"]\n",
     "\n",
-    "agent1 = SampleAgent([GreedyHumanModel(mlam=agent_eval.env.mlam),\n",
-    "                      RandomAgent(all_actions=True),\n",
-    "                      RandomAgent(all_actions=False)])\n",
+    "agent1 = SampleAgent(\n",
+    "    [\n",
+    "        GreedyHumanModel(mlam=agent_eval.env.mlam),\n",
+    "        RandomAgent(all_actions=True),\n",
+    "        RandomAgent(all_actions=False),\n",
+    "    ]\n",
+    ")\n",
     "agent1.agents[0].set_agent_index(0)\n",
-    "agent2 = SampleAgent([GreedyHumanModel(mlam=agent_eval.env.mlam),\n",
-    "                      RandomAgent(all_actions=True),\n",
-    "                      RandomAgent(all_actions=False)])\n",
+    "agent2 = SampleAgent(\n",
+    "    [\n",
+    "        GreedyHumanModel(mlam=agent_eval.env.mlam),\n",
+    "        RandomAgent(all_actions=True),\n",
+    "        RandomAgent(all_actions=False),\n",
+    "    ]\n",
+    ")\n",
     "agent2.agents[0].set_agent_index(1)\n",
     "\n",
     "for i in range(1, 11):\n",
-    "    state = trajectory_random_pair[\"ep_states\"][0][i*100]\n",
-    "    action_probs = [list(agent1.action(state)[1][\"action_probs\"]), list(agent2.action(state)[1][\"action_probs\"])]\n",
-    "    kwargs = {\"hud_data\": {}, \"grid\":grid, \"state\":state.to_dict(), \"action_probs\":action_probs}\n",
-    "    test_dict = {\"config\": config, \"kwargs\": kwargs, \n",
-    "     \"comment\": \"Action probs tests\",\n",
-    "     \"result_array_filename\": \"test_action_probs_display_%d.npy\"%i}\n",
-    "    print(\"test_action_probs_display_%i dict\"%i)\n",
+    "    state = trajectory_random_pair[\"ep_states\"][0][i * 100]\n",
+    "    action_probs = [\n",
+    "        list(agent1.action(state)[1][\"action_probs\"]),\n",
+    "        list(agent2.action(state)[1][\"action_probs\"]),\n",
+    "    ]\n",
+    "    kwargs = {\n",
+    "        \"hud_data\": {},\n",
+    "        \"grid\": grid,\n",
+    "        \"state\": state.to_dict(),\n",
+    "        \"action_probs\": action_probs,\n",
+    "    }\n",
+    "    test_dict = {\n",
+    "        \"config\": config,\n",
+    "        \"kwargs\": kwargs,\n",
+    "        \"comment\": \"Action probs tests\",\n",
+    "        \"result_array_filename\": \"test_action_probs_display_%d.npy\" % i,\n",
+    "    }\n",
+    "    print(\"test_action_probs_display_%i dict\" % i)\n",
     "    test_array = display_and_export_to_array(test_dict)\n",
     "    action_probs_tests.append(test_dict)\n",
     "    action_probs_results.append(test_array)"
@@ -1145,7 +1243,11 @@
    "outputs": [],
    "source": [
     "# if all stuff above is ok run this cell\n",
-    "save_test_files(action_probs_tests, action_probs_results, \"render_state_data_test_action_probs_display.json\")"
+    "save_test_files(\n",
+    "    action_probs_tests,\n",
+    "    action_probs_results,\n",
+    "    \"render_state_data_test_action_probs_display.json\",\n",
+    ")"
    ]
   },
   {
diff --git a/testing/agent_test.py b/testing/agent_test.py
index 79067343..40332104 100644
--- a/testing/agent_test.py
+++ b/testing/agent_test.py
@@ -1,12 +1,27 @@
 import unittest
+
 import numpy as np
 
-from overcooked_ai_py.agents.agent import AgentPair, FixedPlanAgent, GreedyHumanModel, RandomAgent, SampleAgent
-from overcooked_ai_py.mdp.actions import Direction, Action
-from overcooked_ai_py.mdp.overcooked_mdp import OvercookedGridworld, OvercookedState, PlayerState, ObjectState
-from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
-from overcooked_ai_py.planning.planners import MediumLevelActionManager, NO_COUNTERS_PARAMS
+from overcooked_ai_py.agents.agent import (
+    AgentPair,
+    FixedPlanAgent,
+    GreedyHumanModel,
+    RandomAgent,
+    SampleAgent,
+)
 from overcooked_ai_py.agents.benchmarking import AgentEvaluator
+from overcooked_ai_py.mdp.actions import Action, Direction
+from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
+from overcooked_ai_py.mdp.overcooked_mdp import (
+    ObjectState,
+    OvercookedGridworld,
+    OvercookedState,
+    PlayerState,
+)
+from overcooked_ai_py.planning.planners import (
+    NO_COUNTERS_PARAMS,
+    MediumLevelActionManager,
+)
 
 np.random.seed(42)
 
@@ -19,21 +34,26 @@
 force_compute = True
 DISPLAY = False
 
-simple_mdp = OvercookedGridworld.from_layout_name('cramped_room')
-large_mdp = OvercookedGridworld.from_layout_name('corridor')
+simple_mdp = OvercookedGridworld.from_layout_name("cramped_room")
+large_mdp = OvercookedGridworld.from_layout_name("corridor")
 
 
 class TestAgentEvaluator(unittest.TestCase):
-
     def setUp(self):
-        self.agent_eval = AgentEvaluator.from_layout_name({"layout_name": "cramped_room"}, {"horizon": 100})
-        
+        self.agent_eval = AgentEvaluator.from_layout_name(
+            {"layout_name": "cramped_room"}, {"horizon": 100}
+        )
+
     def test_human_model_pair(self):
         trajs = self.agent_eval.evaluate_human_model_pair()
         try:
             AgentEvaluator.check_trajectories(trajs, verbose=False)
         except AssertionError as e:
-            self.fail("Trajectories were not returned in standard format:\n{}".format(e))
+            self.fail(
+                "Trajectories were not returned in standard format:\n{}".format(
+                    e
+                )
+            )
 
     def test_rollouts(self):
         ap = AgentPair(RandomAgent(), RandomAgent())
@@ -41,78 +61,133 @@ def test_rollouts(self):
         try:
             AgentEvaluator.check_trajectories(trajs, verbose=False)
         except AssertionError as e:
-            self.fail("Trajectories were not returned in standard format:\n{}".format(e))
-        
+            self.fail(
+                "Trajectories were not returned in standard format:\n{}".format(
+                    e
+                )
+            )
+
     def test_mlam_computation(self):
         try:
             self.agent_eval.env.mlam
         except Exception as e:
-            self.fail("Failed to compute MediumLevelActionManager:\n{}".format(e))
+            self.fail(
+                "Failed to compute MediumLevelActionManager:\n{}".format(e)
+            )
 
 
 class TestBasicAgents(unittest.TestCase):
-
     def setUp(self):
-        self.mlam_large = MediumLevelActionManager.from_pickle_or_compute(large_mdp, NO_COUNTERS_PARAMS, force_compute=force_compute_large)
+        self.mlam_large = MediumLevelActionManager.from_pickle_or_compute(
+            large_mdp, NO_COUNTERS_PARAMS, force_compute=force_compute_large
+        )
 
     def test_fixed_plan_agents(self):
         a0 = FixedPlanAgent([s, e, n, w])
         a1 = FixedPlanAgent([s, w, n, e])
         agent_pair = AgentPair(a0, a1)
         env = OvercookedEnv.from_mdp(large_mdp, horizon=10)
-        trajectory, time_taken, _, _ = env.run_agents(agent_pair, include_final_state=True, display=DISPLAY)
+        trajectory, time_taken, _, _ = env.run_agents(
+            agent_pair, include_final_state=True, display=DISPLAY
+        )
         end_state = trajectory[-1][0]
         self.assertEqual(time_taken, 10)
-        self.assertEqual(env.mdp.get_standard_start_state().player_positions, end_state.player_positions)
+        self.assertEqual(
+            env.mdp.get_standard_start_state().player_positions,
+            end_state.player_positions,
+        )
 
     def test_two_greedy_human_open_map(self):
-        scenario_2_mdp = OvercookedGridworld.from_layout_name('scenario2')
-        mlam = MediumLevelActionManager.from_pickle_or_compute(scenario_2_mdp, NO_COUNTERS_PARAMS, force_compute=force_compute)
+        scenario_2_mdp = OvercookedGridworld.from_layout_name("scenario2")
+        mlam = MediumLevelActionManager.from_pickle_or_compute(
+            scenario_2_mdp, NO_COUNTERS_PARAMS, force_compute=force_compute
+        )
         a0 = GreedyHumanModel(mlam)
         a1 = GreedyHumanModel(mlam)
         agent_pair = AgentPair(a0, a1)
         start_state = OvercookedState(
-            [P((8, 1), s),
-             P((1, 1), s)],
+            [P((8, 1), s), P((1, 1), s)],
             {},
-            all_orders=scenario_2_mdp.start_all_orders
+            all_orders=scenario_2_mdp.start_all_orders,
+        )
+        env = OvercookedEnv.from_mdp(
+            scenario_2_mdp, start_state_fn=lambda: start_state, horizon=100
+        )
+        trajectory, time_taken, _, _ = env.run_agents(
+            agent_pair, include_final_state=True, display=DISPLAY
         )
-        env = OvercookedEnv.from_mdp(scenario_2_mdp, start_state_fn=lambda: start_state, horizon=100)
-        trajectory, time_taken, _, _ = env.run_agents(agent_pair, include_final_state=True, display=DISPLAY)
 
     def test_sample_agent(self):
-        agent = SampleAgent([RandomAgent(all_actions=False), RandomAgent(all_actions=True)])
+        agent = SampleAgent(
+            [RandomAgent(all_actions=False), RandomAgent(all_actions=True)]
+        )
         probs = agent.action(None)[1]["action_probs"]
-        expected_probs = np.array([0.18333333, 0.18333333, 0.18333333, 0.18333333, 0.18333333, 0.08333333])
+        expected_probs = np.array(
+            [
+                0.18333333,
+                0.18333333,
+                0.18333333,
+                0.18333333,
+                0.18333333,
+                0.08333333,
+            ]
+        )
         self.assertTrue(np.allclose(probs, expected_probs))
 
+
 class TestAgentEvaluatorStatic(unittest.TestCase):
 
-    layout_name_lst = ["asymmetric_advantages", "asymmetric_advantages_tomato", "bonus_order_test", "bottleneck",
-                       "centre_objects", "centre_pots", "corridor", "forced_coordination_tomato", "unident",
-                       "marshmallow_experiment", "marshmallow_experiment_coordination", "you_shall_not_pass"]
+    layout_name_lst = [
+        "asymmetric_advantages",
+        "asymmetric_advantages_tomato",
+        "bonus_order_test",
+        "bottleneck",
+        "centre_objects",
+        "centre_pots",
+        "corridor",
+        "forced_coordination_tomato",
+        "unident",
+        "marshmallow_experiment",
+        "marshmallow_experiment_coordination",
+        "you_shall_not_pass",
+    ]
 
     def test_from_mdp(self):
         for layout_name in self.layout_name_lst:
             orignal_mdp = OvercookedGridworld.from_layout_name(layout_name)
-            ae = AgentEvaluator.from_mdp(mdp=orignal_mdp, env_params={"horizon": 400})
+            ae = AgentEvaluator.from_mdp(
+                mdp=orignal_mdp, env_params={"horizon": 400}
+            )
             ae_mdp = ae.env.mdp
-            self.assertEqual(orignal_mdp, ae_mdp, "mdp with name " + layout_name + " experienced an inconsistency")
+            self.assertEqual(
+                orignal_mdp,
+                ae_mdp,
+                "mdp with name "
+                + layout_name
+                + " experienced an inconsistency",
+            )
 
     def test_from_mdp_params_layout(self):
         for layout_name in self.layout_name_lst:
             orignal_mdp = OvercookedGridworld.from_layout_name(layout_name)
-            ae = AgentEvaluator.from_layout_name(mdp_params={"layout_name": layout_name}, env_params={"horizon": 400})
+            ae = AgentEvaluator.from_layout_name(
+                mdp_params={"layout_name": layout_name},
+                env_params={"horizon": 400},
+            )
             ae_mdp = ae.env.mdp
-            self.assertEqual(orignal_mdp, ae_mdp, "mdp with name " + layout_name + " experienced an inconsistency")
+            self.assertEqual(
+                orignal_mdp,
+                ae_mdp,
+                "mdp with name "
+                + layout_name
+                + " experienced an inconsistency",
+            )
 
     mdp_gen_params_1 = {
         "inner_shape": (10, 7),
         "prop_empty": 0.95,
         "prop_feats": 0.1,
-        "start_all_orders": [
-            {"ingredients": ["onion", "onion", "onion"]}
-        ],
+        "start_all_orders": [{"ingredients": ["onion", "onion", "onion"]}],
         "display": False,
     }
 
@@ -120,9 +195,7 @@ def test_from_mdp_params_layout(self):
         "inner_shape": (10, 7),
         "prop_empty": 0.7,
         "prop_feats": 0.5,
-        "start_all_orders": [
-            {"ingredients": ["onion", "onion", "onion"]}
-        ],
+        "start_all_orders": [{"ingredients": ["onion", "onion", "onion"]}],
         "display": False,
     }
 
@@ -130,9 +203,7 @@ def test_from_mdp_params_layout(self):
         "inner_shape": (10, 7),
         "prop_empty": 0.5,
         "prop_feats": 0.4,
-        "start_all_orders": [
-            {"ingredients": ["onion", "onion", "onion"]}
-        ],
+        "start_all_orders": [{"ingredients": ["onion", "onion", "onion"]}],
         "display": False,
     }
 
@@ -142,56 +213,76 @@ def test_from_mdp_params_layout(self):
 
     def test_from_mdp_params_variable_across(self):
         for mdp_gen_params in self.mdp_gen_params_lst:
-            ae0 = AgentEvaluator.from_mdp_params_infinite(mdp_params=mdp_gen_params,
-                                                          env_params={"horizon": 400, "num_mdp": np.inf},
-                                                          outer_shape=self.outer_shape)
-            ae1 = AgentEvaluator.from_mdp_params_infinite(mdp_params=mdp_gen_params,
-                                                          env_params={"horizon": 400, "num_mdp": np.inf},
-                                                          outer_shape=self.outer_shape)
-            self.assertFalse(ae0.env.mdp == ae1.env.mdp,
-                             "2 randomly generated layouts across 2 evaluators are the same, which is wrong")
+            ae0 = AgentEvaluator.from_mdp_params_infinite(
+                mdp_params=mdp_gen_params,
+                env_params={"horizon": 400, "num_mdp": np.inf},
+                outer_shape=self.outer_shape,
+            )
+            ae1 = AgentEvaluator.from_mdp_params_infinite(
+                mdp_params=mdp_gen_params,
+                env_params={"horizon": 400, "num_mdp": np.inf},
+                outer_shape=self.outer_shape,
+            )
+            self.assertFalse(
+                ae0.env.mdp == ae1.env.mdp,
+                "2 randomly generated layouts across 2 evaluators are the same, which is wrong",
+            )
 
     def test_from_mdp_params_variable_infinite(self):
         for mdp_gen_params in self.mdp_gen_params_lst:
-            ae = AgentEvaluator.from_mdp_params_infinite(mdp_params=mdp_gen_params,
-                                                         env_params={"horizon": 400, "num_mdp": np.inf},
-                                                         outer_shape=self.outer_shape)
+            ae = AgentEvaluator.from_mdp_params_infinite(
+                mdp_params=mdp_gen_params,
+                env_params={"horizon": 400, "num_mdp": np.inf},
+                outer_shape=self.outer_shape,
+            )
             mdp_0 = ae.env.mdp.copy()
             for _ in range(5):
                 ae.env.reset(regen_mdp=True)
                 mdp_1 = ae.env.mdp
-                self.assertFalse(mdp_0 == mdp_1,
-                                 "with infinite layout generator and regen_mdp=True, the 2 layouts should not be the same")
+                self.assertFalse(
+                    mdp_0 == mdp_1,
+                    "with infinite layout generator and regen_mdp=True, the 2 layouts should not be the same",
+                )
 
     def test_from_mdp_params_variable_infinite_no_regen(self):
         for mdp_gen_params in self.mdp_gen_params_lst:
-            ae = AgentEvaluator.from_mdp_params_infinite(mdp_params=mdp_gen_params,
-                                                         env_params={"horizon": 400, "num_mdp": np.inf},
-                                                         outer_shape=self.outer_shape)
+            ae = AgentEvaluator.from_mdp_params_infinite(
+                mdp_params=mdp_gen_params,
+                env_params={"horizon": 400, "num_mdp": np.inf},
+                outer_shape=self.outer_shape,
+            )
             mdp_0 = ae.env.mdp.copy()
             for _ in range(5):
                 ae.env.reset(regen_mdp=False)
                 mdp_1 = ae.env.mdp
-                self.assertTrue(mdp_0 == mdp_1,
-                                 "with infinite layout generator and regen_mdp=False, the 2 layouts should be the same")
+                self.assertTrue(
+                    mdp_0 == mdp_1,
+                    "with infinite layout generator and regen_mdp=False, the 2 layouts should be the same",
+                )
 
     def test_from_mdp_params_variable_infinite_specified(self):
         for mdp_gen_params in self.mdp_gen_params_lst:
-            ae = AgentEvaluator.from_mdp_params_infinite(mdp_params=mdp_gen_params,
-                                                         env_params={"horizon": 400, "num_mdp": np.inf},
-                                                         outer_shape=self.outer_shape)
+            ae = AgentEvaluator.from_mdp_params_infinite(
+                mdp_params=mdp_gen_params,
+                env_params={"horizon": 400, "num_mdp": np.inf},
+                outer_shape=self.outer_shape,
+            )
             mdp_0 = ae.env.mdp.copy()
             for _ in range(5):
                 ae.env.reset(regen_mdp=True)
                 mdp_1 = ae.env.mdp
-                self.assertFalse(mdp_0 == mdp_1,
-                                 "with infinite layout generator and regen_mdp=True, the 2 layouts should not be the same")
+                self.assertFalse(
+                    mdp_0 == mdp_1,
+                    "with infinite layout generator and regen_mdp=True, the 2 layouts should not be the same",
+                )
 
     def test_from_mdp_params_variable_finite(self):
         for mdp_gen_params in self.mdp_gen_params_lst:
-            ae = AgentEvaluator.from_mdp_params_finite(mdp_params=mdp_gen_params,
-                                                       env_params={"horizon": 400, "num_mdp": 2},
-                                                       outer_shape=self.outer_shape)
+            ae = AgentEvaluator.from_mdp_params_finite(
+                mdp_params=mdp_gen_params,
+                env_params={"horizon": 400, "num_mdp": 2},
+                outer_shape=self.outer_shape,
+            )
             mdp_0 = ae.env.mdp.copy()
             seen = [mdp_0]
             for _ in range(20):
@@ -202,18 +293,33 @@ def test_from_mdp_params_variable_finite(self):
                         seen.append(mdp_i.copy())
                 elif len(seen) == 2:
                     mdp_0, mdp_1 = seen
-                    self.assertTrue((mdp_i == mdp_0 or mdp_i == mdp_1),
-                                        "more than 2 mdp was created, the function failed to perform")
+                    self.assertTrue(
+                        (mdp_i == mdp_0 or mdp_i == mdp_1),
+                        "more than 2 mdp was created, the function failed to perform",
+                    )
                 else:
-                    self.assertTrue(False, "theoretically unreachable statement")
-
-    layout_name_short_lst = ["cramped_room", "cramped_room_tomato", "simple_o", "simple_tomato", "simple_o_t"]
+                    self.assertTrue(
+                        False, "theoretically unreachable statement"
+                    )
+
+    layout_name_short_lst = [
+        "cramped_room",
+        "cramped_room_tomato",
+        "simple_o",
+        "simple_tomato",
+        "simple_o_t",
+    ]
     biased = [0.1, 0.15, 0.2, 0.25, 0.3]
     num_reset = 200000
 
     def test_from_mdp_lst_default(self):
-        mdp_lst = [OvercookedGridworld.from_layout_name(name) for name in self.layout_name_short_lst]
-        ae = AgentEvaluator.from_mdp_lst(mdp_lst=mdp_lst, env_params={"horizon": 400})
+        mdp_lst = [
+            OvercookedGridworld.from_layout_name(name)
+            for name in self.layout_name_short_lst
+        ]
+        ae = AgentEvaluator.from_mdp_lst(
+            mdp_lst=mdp_lst, env_params={"horizon": 400}
+        )
         counts = {}
 
         for _ in range(self.num_reset):
@@ -224,11 +330,20 @@ def test_from_mdp_lst_default(self):
                 counts[ae.env.mdp.layout_name] = 1
 
         for k, v in counts.items():
-            self.assertAlmostEqual(0.2, v/self.num_reset, 2, "more than 2 places off for " + k)
+            self.assertAlmostEqual(
+                0.2, v / self.num_reset, 2, "more than 2 places off for " + k
+            )
 
     def test_from_mdp_lst_uniform(self):
-        mdp_lst = [OvercookedGridworld.from_layout_name(name) for name in self.layout_name_short_lst]
-        ae = AgentEvaluator.from_mdp_lst(mdp_lst=mdp_lst, env_params={"horizon": 400}, sampling_freq=[0.2, 0.2, 0.2, 0.2, 0.2])
+        mdp_lst = [
+            OvercookedGridworld.from_layout_name(name)
+            for name in self.layout_name_short_lst
+        ]
+        ae = AgentEvaluator.from_mdp_lst(
+            mdp_lst=mdp_lst,
+            env_params={"horizon": 400},
+            sampling_freq=[0.2, 0.2, 0.2, 0.2, 0.2],
+        )
         counts = {}
 
         for _ in range(self.num_reset):
@@ -239,11 +354,20 @@ def test_from_mdp_lst_uniform(self):
                 counts[ae.env.mdp.layout_name] = 1
 
         for k, v in counts.items():
-            self.assertAlmostEqual(0.2, v/self.num_reset, 2, "more than 2 places off for " + k)
+            self.assertAlmostEqual(
+                0.2, v / self.num_reset, 2, "more than 2 places off for " + k
+            )
 
     def test_from_mdp_lst_biased(self):
-        mdp_lst = [OvercookedGridworld.from_layout_name(name) for name in self.layout_name_short_lst]
-        ae = AgentEvaluator.from_mdp_lst(mdp_lst=mdp_lst, env_params={"horizon": 400}, sampling_freq=self.biased)
+        mdp_lst = [
+            OvercookedGridworld.from_layout_name(name)
+            for name in self.layout_name_short_lst
+        ]
+        ae = AgentEvaluator.from_mdp_lst(
+            mdp_lst=mdp_lst,
+            env_params={"horizon": 400},
+            sampling_freq=self.biased,
+        )
         counts = {}
 
         for _ in range(self.num_reset):
@@ -254,11 +378,16 @@ def test_from_mdp_lst_biased(self):
                 counts[ae.env.mdp.layout_name] = 1
 
         # construct the ground truth
-        gt = {self.layout_name_short_lst[i]: self.biased[i] for i in range(len(self.layout_name_short_lst))}
+        gt = {
+            self.layout_name_short_lst[i]: self.biased[i]
+            for i in range(len(self.layout_name_short_lst))
+        }
 
         for k, v in counts.items():
-            self.assertAlmostEqual(gt[k], v/self.num_reset, 2, "more than 2 places off for " + k)
+            self.assertAlmostEqual(
+                gt[k], v / self.num_reset, 2, "more than 2 places off for " + k
+            )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/testing/mdp_gen_schedule_test.py b/testing/mdp_gen_schedule_test.py
index 7d6f0057..e8d04e94 100644
--- a/testing/mdp_gen_schedule_test.py
+++ b/testing/mdp_gen_schedule_test.py
@@ -1,9 +1,14 @@
 import unittest
+
 import numpy as np
 
-from overcooked_ai_py.mdp.actions import Direction, Action
-from overcooked_ai_py.mdp.overcooked_mdp import OvercookedGridworld, PlayerState, ObjectState
 from overcooked_ai_py.agents.benchmarking import AgentEvaluator
+from overcooked_ai_py.mdp.actions import Action, Direction
+from overcooked_ai_py.mdp.overcooked_mdp import (
+    ObjectState,
+    OvercookedGridworld,
+    PlayerState,
+)
 
 np.random.seed(42)
 
@@ -16,8 +21,8 @@
 force_compute = True
 DISPLAY = False
 
-simple_mdp = OvercookedGridworld.from_layout_name('cramped_room')
-large_mdp = OvercookedGridworld.from_layout_name('corridor')
+simple_mdp = OvercookedGridworld.from_layout_name("cramped_room")
+large_mdp = OvercookedGridworld.from_layout_name("corridor")
 
 
 def params_schedule_fn_constant_09_01(outside_information):
@@ -28,11 +33,9 @@ def params_schedule_fn_constant_09_01(outside_information):
         "inner_shape": (7, 5),
         "prop_empty": 0.9,
         "prop_feats": 0.1,
-        "start_all_orders": [
-            {"ingredients": ["onion", "onion", "onion"]}
-        ],
+        "start_all_orders": [{"ingredients": ["onion", "onion", "onion"]}],
         "display": False,
-        "rew_shaping_params": None
+        "rew_shaping_params": None,
     }
     return mdp_default_gen_params
 
@@ -45,11 +48,9 @@ def params_schedule_fn_constant_07_03(outside_info):
         "inner_shape": (7, 5),
         "prop_empty": 0.7,
         "prop_feats": 0.3,
-        "start_all_orders": [
-            {"ingredients": ["onion", "onion", "onion"]}
-        ],
+        "start_all_orders": [{"ingredients": ["onion", "onion", "onion"]}],
         "display": False,
-        "rew_shaping_params": None
+        "rew_shaping_params": None,
     }
     return mdp_default_gen_params
 
@@ -62,11 +63,9 @@ def params_schedule_fn_constant_05_05(outside_info):
         "inner_shape": (7, 5),
         "prop_empty": 0.5,
         "prop_feats": 0.5,
-        "start_all_orders": [
-            {"ingredients": ["onion", "onion", "onion"]}
-        ],
+        "start_all_orders": [{"ingredients": ["onion", "onion", "onion"]}],
         "display": False,
-        "rew_shaping_params": None
+        "rew_shaping_params": None,
     }
     return mdp_default_gen_params
 
@@ -76,8 +75,9 @@ def params_schedule_fn_interval(outside_info):
     outside_information (dict):
         progress (float in [0, 1] interval) a number that indicate progress
     """
-    assert outside_info != {} and "progress" in outside_info, \
-        "if this happens during initialization, please add initial_info to env_params to address the issue"
+    assert (
+        outside_info != {} and "progress" in outside_info
+    ), "if this happens during initialization, please add initial_info to env_params to address the issue"
     progress = outside_info["progress"]
     prop_empty = 0.9 - 0.4 * progress
     prop_feats = 0.1 + 0.4 * progress
@@ -85,61 +85,71 @@ def params_schedule_fn_interval(outside_info):
         "inner_shape": (7, 5),
         "prop_empty": prop_empty,
         "prop_feats": prop_feats,
-        "start_all_orders": [
-            {"ingredients": ["onion", "onion", "onion"]}
-        ],
+        "start_all_orders": [{"ingredients": ["onion", "onion", "onion"]}],
         "display": False,
-        "rew_shaping_params": None
+        "rew_shaping_params": None,
     }
     return mdp_params_generated
 
 
 default_env_params_infinite = {"horizon": 400, "num_mdp": np.inf}
 
-default_env_params_infinite_interval = {"horizon": 400, "num_mdp": np.inf, "initial_info": {"progress": 0}}
+default_env_params_infinite_interval = {
+    "horizon": 400,
+    "num_mdp": np.inf,
+    "initial_info": {"progress": 0},
+}
 
 
 class TestParamScheduleFnConstant(unittest.TestCase):
-
     def test_constant_schedule_095_01(self):
-        ae = AgentEvaluator.from_mdp_params_infinite(mdp_params=None, env_params=default_env_params_infinite,
-                                                     outer_shape= (7, 5),
-                                                     mdp_params_schedule_fn=params_schedule_fn_constant_09_01)
+        ae = AgentEvaluator.from_mdp_params_infinite(
+            mdp_params=None,
+            env_params=default_env_params_infinite,
+            outer_shape=(7, 5),
+            mdp_params_schedule_fn=params_schedule_fn_constant_09_01,
+        )
         num_empty_grid = []
         for i in range(500):
             ae.env.reset()
-            empty_i = len(ae.env.mdp.terrain_pos_dict[' '])
+            empty_i = len(ae.env.mdp.terrain_pos_dict[" "])
             num_empty_grid.append(empty_i)
-        avg_num_empty = sum(num_empty_grid)/len(num_empty_grid)
+        avg_num_empty = sum(num_empty_grid) / len(num_empty_grid)
         print("avg number of empty grid:", avg_num_empty)
         # the number of empty square should be consistant"
 
         self.assertTrue(13.9 < avg_num_empty < 14.1)
 
     def test_constant_schedule_07_03(self):
-        ae = AgentEvaluator.from_mdp_params_infinite(mdp_params=None, env_params=default_env_params_infinite,
-                                                     outer_shape= (7, 5),
-                                                     mdp_params_schedule_fn=params_schedule_fn_constant_07_03)
+        ae = AgentEvaluator.from_mdp_params_infinite(
+            mdp_params=None,
+            env_params=default_env_params_infinite,
+            outer_shape=(7, 5),
+            mdp_params_schedule_fn=params_schedule_fn_constant_07_03,
+        )
         num_empty_grid = []
         for i in range(500):
             ae.env.reset()
-            empty_i = len(ae.env.mdp.terrain_pos_dict[' '])
+            empty_i = len(ae.env.mdp.terrain_pos_dict[" "])
             num_empty_grid.append(empty_i)
-        avg_num_empty = sum(num_empty_grid)/len(num_empty_grid)
+        avg_num_empty = sum(num_empty_grid) / len(num_empty_grid)
         print("avg number of empty grid:", avg_num_empty)
         # the number of empty square should be fairlyconsistant"
         self.assertTrue(11.5 < avg_num_empty < 11.8)
 
     def test_constant_schedule_05_05(self):
-        ae = AgentEvaluator.from_mdp_params_infinite(mdp_params=None, env_params=default_env_params_infinite,
-                                                     outer_shape= (7, 5),
-                                                     mdp_params_schedule_fn=params_schedule_fn_constant_05_05)
+        ae = AgentEvaluator.from_mdp_params_infinite(
+            mdp_params=None,
+            env_params=default_env_params_infinite,
+            outer_shape=(7, 5),
+            mdp_params_schedule_fn=params_schedule_fn_constant_05_05,
+        )
         num_empty_grid = []
         for i in range(500):
             ae.env.reset()
-            empty_i = len(ae.env.mdp.terrain_pos_dict[' '])
+            empty_i = len(ae.env.mdp.terrain_pos_dict[" "])
             num_empty_grid.append(empty_i)
-        avg_num_empty = sum(num_empty_grid)/len(num_empty_grid)
+        avg_num_empty = sum(num_empty_grid) / len(num_empty_grid)
         print("avg number of empty grid:", avg_num_empty)
         # the number of empty square should be fairlyconsistant"
         self.assertTrue(10.4 < avg_num_empty < 10.9)
@@ -147,13 +157,16 @@ def test_constant_schedule_05_05(self):
 
 class TestParamScheduleFnInterval(unittest.TestCase):
     def test_interval_schedule(self):
-        ae = AgentEvaluator.from_mdp_params_infinite(mdp_params=None, env_params=default_env_params_infinite_interval,
-                                                     outer_shape= (7, 5),
-                                                     mdp_params_schedule_fn=params_schedule_fn_interval)
+        ae = AgentEvaluator.from_mdp_params_infinite(
+            mdp_params=None,
+            env_params=default_env_params_infinite_interval,
+            outer_shape=(7, 5),
+            mdp_params_schedule_fn=params_schedule_fn_interval,
+        )
         num_empty_grid = []
         for i in range(4000):
-            ae.env.reset(outside_info={"progress": i/4000})
-            empty_i = len(ae.env.mdp.terrain_pos_dict[' '])
+            ae.env.reset(outside_info={"progress": i / 4000})
+            empty_i = len(ae.env.mdp.terrain_pos_dict[" "])
             num_empty_grid.append(empty_i)
         avg_num_empty_09_01 = sum(num_empty_grid[0:50]) / 50
         self.assertTrue(13.9 < avg_num_empty_09_01 < 14.1)
@@ -161,9 +174,13 @@ def test_interval_schedule(self):
         self.assertTrue(11.5 < avg_num_empty_07_03 < 11.8)
         avg_num_empty_05_05 = sum(num_empty_grid[3950:4000]) / 50
         self.assertTrue(10.4 < avg_num_empty_05_05 < 10.9)
-        print("avg number of empty grids:", avg_num_empty_09_01, avg_num_empty_07_03, avg_num_empty_05_05)
+        print(
+            "avg number of empty grids:",
+            avg_num_empty_09_01,
+            avg_num_empty_07_03,
+            avg_num_empty_05_05,
+        )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
-
diff --git a/testing/overcooked_test.py b/testing/overcooked_test.py
index fb0ebd8b..aa46b248 100644
--- a/testing/overcooked_test.py
+++ b/testing/overcooked_test.py
@@ -1,18 +1,62 @@
-import unittest, os, shutil, glob
-import json, copy
-import numpy as np
+import copy
+import glob
+import json
+import os
+import shutil
+import unittest
 from math import factorial
-from overcooked_ai_py.mdp.actions import Action, Direction
-from overcooked_ai_py.mdp.overcooked_mdp import PlayerState, OvercookedGridworld, OvercookedState, ObjectState, SoupState, Recipe
-from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv, DEFAULT_ENV_PARAMS
-from overcooked_ai_py.mdp.overcooked_trajectory import append_trajectories, DEFAULT_TRAJ_KEYS, TIMESTEP_TRAJ_KEYS, EPISODE_TRAJ_KEYS
-from overcooked_ai_py.mdp.layout_generator import LayoutGenerator, ONION_DISPENSER, TOMATO_DISPENSER, POT, DISH_DISPENSER, SERVING_LOC
-from overcooked_ai_py.agents.agent import AgentGroup, AgentPair, GreedyHumanModel, FixedPlanAgent, RandomAgent
+
+import numpy as np
+from utils import generate_serialized_trajectory
+
+from overcooked_ai_py.agents.agent import (
+    AgentGroup,
+    AgentPair,
+    FixedPlanAgent,
+    GreedyHumanModel,
+    RandomAgent,
+)
 from overcooked_ai_py.agents.benchmarking import AgentEvaluator
-from overcooked_ai_py.planning.planners import MediumLevelActionManager, NO_COUNTERS_PARAMS, MotionPlanner
-from overcooked_ai_py.utils import save_pickle, load_pickle, iterate_over_json_files_in_dir, load_from_json, save_as_json
+from overcooked_ai_py.mdp.actions import Action, Direction
+from overcooked_ai_py.mdp.layout_generator import (
+    DISH_DISPENSER,
+    ONION_DISPENSER,
+    POT,
+    SERVING_LOC,
+    TOMATO_DISPENSER,
+    LayoutGenerator,
+)
+from overcooked_ai_py.mdp.overcooked_env import (
+    DEFAULT_ENV_PARAMS,
+    OvercookedEnv,
+)
+from overcooked_ai_py.mdp.overcooked_mdp import (
+    ObjectState,
+    OvercookedGridworld,
+    OvercookedState,
+    PlayerState,
+    Recipe,
+    SoupState,
+)
+from overcooked_ai_py.mdp.overcooked_trajectory import (
+    DEFAULT_TRAJ_KEYS,
+    EPISODE_TRAJ_KEYS,
+    TIMESTEP_TRAJ_KEYS,
+    append_trajectories,
+)
+from overcooked_ai_py.planning.planners import (
+    NO_COUNTERS_PARAMS,
+    MediumLevelActionManager,
+    MotionPlanner,
+)
 from overcooked_ai_py.static import TESTING_DATA_DIR
-from utils import generate_serialized_trajectory
+from overcooked_ai_py.utils import (
+    iterate_over_json_files_in_dir,
+    load_from_json,
+    load_pickle,
+    save_as_json,
+    save_pickle,
+)
 
 START_ORDER_LIST = ["any"]
 n, s = Direction.NORTH, Direction.SOUTH
@@ -20,12 +64,12 @@
 stay, interact = Action.STAY, Action.INTERACT
 P, Obj = PlayerState, ObjectState
 
+
 def comb(n, k):
     return factorial(n) / (factorial(n - k) * factorial(k))
 
 
 class TestRecipe(unittest.TestCase):
-
     def setUp(self):
         Recipe.configure({})
         self.r1 = Recipe([Recipe.ONION, Recipe.ONION, Recipe.ONION])
@@ -37,7 +81,7 @@ def setUp(self):
 
         self.recipes = [self.r1, self.r2, self.r3, self.r4, self.r5, self.r6]
 
-        self.pickle_temp_dir = os.path.join(TESTING_DATA_DIR, 'recipes')
+        self.pickle_temp_dir = os.path.join(TESTING_DATA_DIR, "recipes")
 
         if not os.path.exists(self.pickle_temp_dir):
             os.makedirs(self.pickle_temp_dir)
@@ -47,13 +91,14 @@ def tearDown(self):
 
         if os.path.exists(self.pickle_temp_dir):
             shutil.rmtree(self.pickle_temp_dir)
-            
 
     def test_eq(self):
 
         self.assertEqual(self.r1, self.r2, "Failed basic equality check")
         self.assertNotEqual(self.r1, self.r3, "Failed Basic inequality check")
-        self.assertNotEqual(self.r1, self.r6, "Failed inequality check with all one ingredient")
+        self.assertNotEqual(
+            self.r1, self.r6, "Failed inequality check with all one ingredient"
+        )
         self.assertEqual(self.r3, self.r4, "Failed basic equality check")
         self.assertEqual(self.r4, self.r5, "Failed ordered equality check")
 
@@ -62,23 +107,26 @@ def test_caching(self):
         self.assertIs(self.r1, self.r2)
         self.assertIs(self.r3, self.r4)
         self.assertIs(self.r4, self.r5)
-        self.assertFalse(self.r6 is self.r1, "different recipes cached to same value")
+        self.assertFalse(
+            self.r6 is self.r1, "different recipes cached to same value"
+        )
 
     def test_serialization(self):
         loaded_recipes = []
 
         # Save and then load every recipe instance
         for i, recipe in enumerate(self.recipes):
-            pickle_path = os.path.join(self.pickle_temp_dir, 'recipe_{}'.format(i))
+            pickle_path = os.path.join(
+                self.pickle_temp_dir, "recipe_{}".format(i)
+            )
             save_pickle(recipe, pickle_path)
             loaded = load_pickle(pickle_path)
             loaded_recipes.append(loaded)
-        
+
         # Ensure loaded recipes equal corresponding original recipe
         for original, loaded in zip(self.recipes, loaded_recipes):
             self.assertEqual(original, loaded)
 
-
     def test_value(self):
         # TODO
         for recipe in self.recipes:
@@ -93,49 +141,112 @@ def test_all_recipes(self):
         for recipe in self.recipes:
             self.assertTrue(recipe in Recipe.ALL_RECIPES)
 
-        self.assertEqual(len(Recipe.ALL_RECIPES), self._expected_num_recipes(len(Recipe.ALL_INGREDIENTS), Recipe.MAX_NUM_INGREDIENTS))
+        self.assertEqual(
+            len(Recipe.ALL_RECIPES),
+            self._expected_num_recipes(
+                len(Recipe.ALL_INGREDIENTS), Recipe.MAX_NUM_INGREDIENTS
+            ),
+        )
 
-        Recipe.configure({ "max_num_ingredients" : 4 })
+        Recipe.configure({"max_num_ingredients": 4})
 
-        self.assertEqual(len(Recipe.ALL_RECIPES), self._expected_num_recipes(len(Recipe.ALL_INGREDIENTS), 4))
+        self.assertEqual(
+            len(Recipe.ALL_RECIPES),
+            self._expected_num_recipes(len(Recipe.ALL_INGREDIENTS), 4),
+        )
 
     def test_invalid_input(self):
 
-        self.assertRaises(ValueError, Recipe, [Recipe.ONION, Recipe.TOMATO, "carrot"])
-        self.assertRaises(ValueError, Recipe, [Recipe.ONION]*4)
+        self.assertRaises(
+            ValueError, Recipe, [Recipe.ONION, Recipe.TOMATO, "carrot"]
+        )
+        self.assertRaises(ValueError, Recipe, [Recipe.ONION] * 4)
         self.assertRaises(ValueError, Recipe, [])
         self.assertRaises(ValueError, Recipe, "invalid argument")
 
     def test_recipes_generation(self):
-        self.assertRaises(AssertionError, Recipe.generate_random_recipes, max_size=Recipe.MAX_NUM_INGREDIENTS+1)
-        self.assertRaises(AssertionError, Recipe.generate_random_recipes, min_size=0)
-        self.assertRaises(AssertionError, Recipe.generate_random_recipes, min_size=3, max_size=2)
-        self.assertRaises(AssertionError, Recipe.generate_random_recipes, ingredients=["onion", "tomato", "fake_ingredient"])
-        self.assertRaises(AssertionError, Recipe.generate_random_recipes, n=99999)
+        self.assertRaises(
+            AssertionError,
+            Recipe.generate_random_recipes,
+            max_size=Recipe.MAX_NUM_INGREDIENTS + 1,
+        )
+        self.assertRaises(
+            AssertionError, Recipe.generate_random_recipes, min_size=0
+        )
+        self.assertRaises(
+            AssertionError,
+            Recipe.generate_random_recipes,
+            min_size=3,
+            max_size=2,
+        )
+        self.assertRaises(
+            AssertionError,
+            Recipe.generate_random_recipes,
+            ingredients=["onion", "tomato", "fake_ingredient"],
+        )
+        self.assertRaises(
+            AssertionError, Recipe.generate_random_recipes, n=99999
+        )
         self.assertEqual(len(Recipe.generate_random_recipes(n=3)), 3)
-        self.assertEqual(len(Recipe.generate_random_recipes(n=99, unique=False)), 99)
+        self.assertEqual(
+            len(Recipe.generate_random_recipes(n=99, unique=False)), 99
+        )
 
-        two_sized_recipes = [Recipe(["onion", "onion"]), Recipe(["onion", "tomato"]), Recipe(["tomato", "tomato"])]
+        two_sized_recipes = [
+            Recipe(["onion", "onion"]),
+            Recipe(["onion", "tomato"]),
+            Recipe(["tomato", "tomato"]),
+        ]
         for _ in range(100):
-            self.assertCountEqual(two_sized_recipes, Recipe.generate_random_recipes(n=3, min_size=2, max_size=2, ingredients=["onion", "tomato"]))
-
-        only_onions_recipes = [Recipe(["onion", "onion"]), Recipe(["onion", "onion", "onion"])]
+            self.assertCountEqual(
+                two_sized_recipes,
+                Recipe.generate_random_recipes(
+                    n=3,
+                    min_size=2,
+                    max_size=2,
+                    ingredients=["onion", "tomato"],
+                ),
+            )
+
+        only_onions_recipes = [
+            Recipe(["onion", "onion"]),
+            Recipe(["onion", "onion", "onion"]),
+        ]
         for _ in range(100):
-            self.assertCountEqual(only_onions_recipes, Recipe.generate_random_recipes(n=2, min_size=2, max_size=3, ingredients=["onion"]))
-        
-        self.assertCountEqual(only_onions_recipes, set([Recipe.generate_random_recipes(n=1, recipes=only_onions_recipes)[0] for _ in range(100)])) # false positives rate for this test is 1/10^99 
+            self.assertCountEqual(
+                only_onions_recipes,
+                Recipe.generate_random_recipes(
+                    n=2, min_size=2, max_size=3, ingredients=["onion"]
+                ),
+            )
+
+        self.assertCountEqual(
+            only_onions_recipes,
+            set(
+                [
+                    Recipe.generate_random_recipes(
+                        n=1, recipes=only_onions_recipes
+                    )[0]
+                    for _ in range(100)
+                ]
+            ),
+        )  # false positives rate for this test is 1/10^99
 
     def _expected_num_recipes(self, num_ingredients, max_len):
         return comb(num_ingredients + max_len, num_ingredients) - 1
 
+
 class TestSoupState(unittest.TestCase):
-    
     def setUp(self):
         Recipe.configure({})
         self.s1 = SoupState.get_soup((0, 0), num_onions=0, num_tomatoes=0)
         self.s2 = SoupState.get_soup((0, 1), num_onions=2, num_tomatoes=1)
-        self.s3 = SoupState.get_soup((1, 1), num_onions=1, num_tomatoes=0, cooking_tick=1)
-        self.s4 = SoupState.get_soup((1, 0), num_onions=0, num_tomatoes=2, finished=True)
+        self.s3 = SoupState.get_soup(
+            (1, 1), num_onions=1, num_tomatoes=0, cooking_tick=1
+        )
+        self.s4 = SoupState.get_soup(
+            (1, 0), num_onions=0, num_tomatoes=2, finished=True
+        )
 
     def test_position(self):
         new_pos = (2, 0)
@@ -172,11 +283,11 @@ def test_is_full(self):
     def test_cooking(self):
         self.s1.add_ingredient_from_str(Recipe.ONION)
         self.s1.add_ingredient_from_str(Recipe.TOMATO)
-        
+
         self.assertTrue(self.s1.is_idle)
         self.assertFalse(self.s1.is_cooking)
         self.assertFalse(self.s1.is_full)
-        
+
         self.s1.begin_cooking()
 
         self.assertFalse(self.s1.is_idle)
@@ -193,30 +304,40 @@ def test_cooking(self):
 
     def test_attributes(self):
         self.assertListEqual(self.s1.ingredients, [])
-        self.assertListEqual(self.s2.ingredients, [Recipe.ONION, Recipe.ONION, Recipe.TOMATO])
+        self.assertListEqual(
+            self.s2.ingredients, [Recipe.ONION, Recipe.ONION, Recipe.TOMATO]
+        )
         self.assertListEqual(self.s3.ingredients, [Recipe.ONION])
-        self.assertListEqual(self.s4.ingredients, [Recipe.TOMATO, Recipe.TOMATO])
+        self.assertListEqual(
+            self.s4.ingredients, [Recipe.TOMATO, Recipe.TOMATO]
+        )
 
         try:
             self.s1.recipe
             self.fail("Expected ValueError to be raised")
-        except ValueError as e: 
+        except ValueError as e:
             pass
         except Exception as e:
-            self.fail("Expected ValueError to be raised, {} raised instead".format(e))
+            self.fail(
+                "Expected ValueError to be raised, {} raised instead".format(e)
+            )
 
         try:
             self.s2.recipe
             self.fail("Expected ValueError to be raised")
-        except ValueError as e: 
+        except ValueError as e:
             pass
         except Exception as e:
-            self.fail("Expected ValueError to be raised, {} raised instead".format(e))
+            self.fail(
+                "Expected ValueError to be raised, {} raised instead".format(e)
+            )
         self.assertEqual(self.s3.recipe, Recipe([Recipe.ONION]))
-        self.assertEqual(self.s4.recipe, Recipe([Recipe.TOMATO, Recipe.TOMATO]))
+        self.assertEqual(
+            self.s4.recipe, Recipe([Recipe.TOMATO, Recipe.TOMATO])
+        )
 
     def test_invalid_ops(self):
-        
+
         # Cannot cook an empty soup
         self.assertRaises(ValueError, self.s1.begin_cooking)
 
@@ -233,10 +354,14 @@ def test_invalid_ops(self):
         self.assertRaises(ValueError, self.s4.begin_cooking)
 
         # Cannot add ingredients to a soup that is cooking
-        self.assertRaises(ValueError, self.s3.add_ingredient_from_str, Recipe.ONION)
+        self.assertRaises(
+            ValueError, self.s3.add_ingredient_from_str, Recipe.ONION
+        )
 
         # Cannot add ingredients to a soup that is ready
-        self.assertRaises(ValueError, self.s4.add_ingredient_from_str, Recipe.ONION)
+        self.assertRaises(
+            ValueError, self.s4.add_ingredient_from_str, Recipe.ONION
+        )
 
         # Cannot remove an ingredient from a soup that is ready
         self.assertRaises(ValueError, self.s4.pop_ingredient)
@@ -248,9 +373,7 @@ def test_invalid_ops(self):
         self.assertRaises(ValueError, self.s1.pop_ingredient)
 
 
-
 class TestDirection(unittest.TestCase):
-
     def test_direction_number_conversion(self):
         all_directions = Direction.ALL_DIRECTIONS
         all_numbers = []
@@ -264,7 +387,7 @@ def test_direction_number_conversion(self):
         # Check that all directions are distinct
         num_directions = len(all_directions)
         self.assertEqual(len(set(all_directions)), num_directions)
-        
+
         # Check that the numbers are 0, 1, ... num_directions - 1
         self.assertEqual(set(all_numbers), set(range(num_directions)))
 
@@ -278,82 +401,114 @@ class TestGridworld(unittest.TestCase):
     def setUp(self):
         self.base_mdp = OvercookedGridworld.from_layout_name("mdp_test")
 
-
     def test_constructor_invalid_inputs(self):
         # Height and width must be at least 3.
         with self.assertRaises(AssertionError):
-            mdp = OvercookedGridworld.from_grid(['X', 'X', 'X'])
+            mdp = OvercookedGridworld.from_grid(["X", "X", "X"])
         with self.assertRaises(AssertionError):
-            mdp = OvercookedGridworld.from_grid([['X', 'X', 'X']])
+            mdp = OvercookedGridworld.from_grid([["X", "X", "X"]])
         with self.assertRaises(AssertionError):
             # Borders must be present.
-            mdp = OvercookedGridworld.from_grid(['XOSX',
-                                                 'P  D',
-                                                 ' 21 '])
+            mdp = OvercookedGridworld.from_grid(["XOSX", "P  D", " 21 "])
 
         with self.assertRaises(AssertionError):
             # The grid can't be ragged.
-            mdp = OvercookedGridworld.from_grid(['XXPXX',
-                                                 'O  2XX',
-                                                 'X1 3 X',
-                                                 'XDXSXX'])
+            mdp = OvercookedGridworld.from_grid(
+                ["XXPXX", "O  2XX", "X1 3 X", "XDXSXX"]
+            )
 
         with self.assertRaises(AssertionError):
             # The agents must be numbered 1 and 2.
-            mdp = OvercookedGridworld.from_grid(['XXPXX',
-                                                 'O  3O',
-                                                 'X1  X',
-                                                 'XDXSX'])
+            mdp = OvercookedGridworld.from_grid(
+                ["XXPXX", "O  3O", "X1  X", "XDXSX"]
+            )
 
         with self.assertRaises(AssertionError):
             # The agents must be numbered 1 and 2.
-            mdp = OvercookedGridworld.from_grid(['XXPXX',
-                                                 'O  1O',
-                                                 'X1  X',
-                                                 'XDXSX'])
+            mdp = OvercookedGridworld.from_grid(
+                ["XXPXX", "O  1O", "X1  X", "XDXSX"]
+            )
 
         with self.assertRaises(AssertionError):
             # B is not a valid element.
-            mdp = OvercookedGridworld.from_grid(['XBPXX',
-                                                 'O  2O',
-                                                 'X1  X',
-                                                 'XDXSX'])
+            mdp = OvercookedGridworld.from_grid(
+                ["XBPXX", "O  2O", "X1  X", "XDXSX"]
+            )
 
     def test_start_positions(self):
         actual_start_state = self.base_mdp.get_standard_start_state()
-        expected_state_path = os.path.join(TESTING_DATA_DIR, "test_start_positions", "expected.json")
+        expected_state_path = os.path.join(
+            TESTING_DATA_DIR, "test_start_positions", "expected.json"
+        )
 
         # NOTE: Uncomment the following line if expected start state deliberately changed
         # save_as_json(actual_start_state.to_dict(), expected_state_path)
 
-        expected_start_state = OvercookedState.from_dict(load_from_json(expected_state_path))
-        self.assertEqual(actual_start_state, expected_start_state, '\n' + str(actual_start_state) + '\n' + str(expected_start_state))
+        expected_start_state = OvercookedState.from_dict(
+            load_from_json(expected_state_path)
+        )
+        self.assertEqual(
+            actual_start_state,
+            expected_start_state,
+            "\n" + str(actual_start_state) + "\n" + str(expected_start_state),
+        )
 
     def test_file_constructor(self):
-        mdp = OvercookedGridworld.from_layout_name('corridor')
+        mdp = OvercookedGridworld.from_layout_name("corridor")
         expected_start_state = OvercookedState(
-            [PlayerState((3, 1), Direction.NORTH), PlayerState((10, 1), Direction.NORTH)], {},
-            all_orders=[{ "ingredients" : ["onion", "onion", "onion"]}])
+            [
+                PlayerState((3, 1), Direction.NORTH),
+                PlayerState((10, 1), Direction.NORTH),
+            ],
+            {},
+            all_orders=[{"ingredients": ["onion", "onion", "onion"]}],
+        )
         actual_start_state = mdp.get_standard_start_state()
-        self.assertEqual(actual_start_state, expected_start_state, '\n' + str(actual_start_state) + '\n' + str(expected_start_state))
+        self.assertEqual(
+            actual_start_state,
+            expected_start_state,
+            "\n" + str(actual_start_state) + "\n" + str(expected_start_state),
+        )
 
     def test_actions(self):
         bad_state = OvercookedState(
-            [PlayerState((0, 0), Direction.NORTH), PlayerState((3, 1), Direction.NORTH)], {})
+            [
+                PlayerState((0, 0), Direction.NORTH),
+                PlayerState((3, 1), Direction.NORTH),
+            ],
+            {},
+        )
         with self.assertRaises(AssertionError):
             self.base_mdp.get_actions(bad_state)
 
-        self.assertEqual(self.base_mdp.get_actions(self.base_mdp.get_standard_start_state()),
-                         [Action.ALL_ACTIONS, Action.ALL_ACTIONS])
+        self.assertEqual(
+            self.base_mdp.get_actions(
+                self.base_mdp.get_standard_start_state()
+            ),
+            [Action.ALL_ACTIONS, Action.ALL_ACTIONS],
+        )
 
     def test_from_dict(self):
-        state_dict = {"players": [{"position": [2, 1], "orientation": [0, -1], "held_object": None }, {"position": [1, 1], "orientation": [0, -1], "held_object": None }], "objects": [{"name": "onion", "position": [1, 0], "state": None }], "order_list": None }
+        state_dict = {
+            "players": [
+                {
+                    "position": [2, 1],
+                    "orientation": [0, -1],
+                    "held_object": None,
+                },
+                {
+                    "position": [1, 1],
+                    "orientation": [0, -1],
+                    "held_object": None,
+                },
+            ],
+            "objects": [{"name": "onion", "position": [1, 0], "state": None}],
+            "order_list": None,
+        }
         state = OvercookedState.from_dict(state_dict)
 
-
     def test_transitions_and_environment(self):
-        bad_state = OvercookedState(
-            [P((0, 0), s), P((3, 1), s)], {})
+        bad_state = OvercookedState([P((0, 0), s), P((3, 1), s)], {})
 
         with self.assertRaises(AssertionError):
             self.base_mdp.get_state_transition(bad_state, stay)
@@ -365,46 +520,65 @@ def check_transition(action, expected_path, recompute=False):
             state = env.state
             pred_state, _ = self.base_mdp.get_state_transition(state, action)
             new_state, sparse_reward, _, _ = env.step(action)
-            self.assertEqual(pred_state, new_state, '\n' + str(pred_state) + '\n' + str(new_state))
+            self.assertEqual(
+                pred_state,
+                new_state,
+                "\n" + str(pred_state) + "\n" + str(new_state),
+            )
 
             # Recompute expected values if desired
             if recompute:
                 actual = {
-                    "state" : pred_state.to_dict(),
-                    "reward" : sparse_reward
+                    "state": pred_state.to_dict(),
+                    "reward": sparse_reward,
                 }
                 save_as_json(actual, expected_path)
 
             # Compute expected values
             expected = load_from_json(expected_path)
-            expected_state = OvercookedState.from_dict(expected['state'])
-            expected_reward = expected['reward']
-            
+            expected_state = OvercookedState.from_dict(expected["state"])
+            expected_reward = expected["reward"]
+
             # Make sure everything lines up (note __eq__ is transitive)
-            self.assertTrue(pred_state.time_independent_equal(expected_state), '\n' + str(pred_state) + '\n' + str(expected_state))
+            self.assertTrue(
+                pred_state.time_independent_equal(expected_state),
+                "\n" + str(pred_state) + "\n" + str(expected_state),
+            )
             self.assertEqual(sparse_reward, expected_reward)
 
-        expected_path = os.path.join(TESTING_DATA_DIR, "test_transitions_and_environments", "expected.json")
+        expected_path = os.path.join(
+            TESTING_DATA_DIR,
+            "test_transitions_and_environments",
+            "expected.json",
+        )
 
         # NOTE: set 'recompute=True' if deliberately updating state dynamics
         check_transition([n, e], expected_path, recompute=False)
 
     def test_mdp_dynamics(self):
-        traj_path = os.path.join(TESTING_DATA_DIR, 'test_mdp_dynamics', 'expected.json')
+        traj_path = os.path.join(
+            TESTING_DATA_DIR, "test_mdp_dynamics", "expected.json"
+        )
 
         # NOTE: uncomment the following line to recompute trajectories if MDP dymamics were deliberately updated
         # generate_serialized_trajectory(self.base_mdp, traj_path)
 
         test_trajectory = AgentEvaluator.load_traj_from_json(traj_path)
-        AgentEvaluator.check_trajectories(test_trajectory, from_json=True, verbose=False)
+        AgentEvaluator.check_trajectories(
+            test_trajectory, from_json=True, verbose=False
+        )
 
     def test_mdp_serialization(self):
         # Where to store serialized states -- will be overwritten each timestep
-        dummy_path = os.path.join(TESTING_DATA_DIR, 'test_mdp_serialization', 'dummy.json')
+        dummy_path = os.path.join(
+            TESTING_DATA_DIR, "test_mdp_serialization", "dummy.json"
+        )
 
         # Get starting seed and random agent pair
         seed = 47
-        random_pair = AgentPair(RandomAgent(all_actions=True), RandomAgent(all_actions=True))
+        random_pair = AgentPair(
+            RandomAgent(all_actions=True), RandomAgent(all_actions=True)
+        )
 
         # Run rollouts with different seeds until sparse reward is achieved
         sparse_reward = 0
@@ -413,13 +587,23 @@ def test_mdp_serialization(self):
             state = self.base_mdp.get_standard_start_state()
             for _ in range(1500):
                 # Ensure serialization and deserializations are inverses
-                reconstructed_state = OvercookedState.from_dict(load_from_json(save_as_json(state.to_dict(), dummy_path)))
-                self.assertEqual(state, reconstructed_state, "\nState: \t\t\t{}\nReconstructed State: \t{}".format(state, reconstructed_state))
+                reconstructed_state = OvercookedState.from_dict(
+                    load_from_json(save_as_json(state.to_dict(), dummy_path))
+                )
+                self.assertEqual(
+                    state,
+                    reconstructed_state,
+                    "\nState: \t\t\t{}\nReconstructed State: \t{}".format(
+                        state, reconstructed_state
+                    ),
+                )
 
                 # Advance state
                 joint_action, _ = zip(*random_pair.joint_action(state))
-                state, infos = self.base_mdp.get_state_transition(state, joint_action)
-                sparse_reward += sum(infos['sparse_reward_by_agent'])
+                state, infos = self.base_mdp.get_state_transition(
+                    state, joint_action
+                )
+                sparse_reward += sum(infos["sparse_reward_by_agent"])
             seed += 1
 
     def test_four_player_mdp(self):
@@ -440,47 +624,71 @@ def test_potential_function(self):
             print("potential: ", self.base_mdp.potential_function(state, mp))
         actions = [Direction.EAST, Action.INTERACT]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [Action.STAY, action])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [Action.STAY, action]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val1 = self.base_mdp.potential_function(state, mp)
-        
 
         # Pick up tomato
         if self.verbose:
             print("pick up tomtato")
         actions = [Direction.WEST, Action.INTERACT]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val2 = self.base_mdp.potential_function(state, mp)
 
-        self.assertLess(val0, val1, "Picking up onion should increase potential")
-        self.assertLess(val1, val2, "Picking up tomato should increase potential")
+        self.assertLess(
+            val0, val1, "Picking up onion should increase potential"
+        )
+        self.assertLess(
+            val1, val2, "Picking up tomato should increase potential"
+        )
 
         # Pot tomato
         if self.verbose:
             print("pot tomato")
-        actions = [Direction.EAST, Direction.NORTH, Action.INTERACT, Direction.WEST]
+        actions = [
+            Direction.EAST,
+            Direction.NORTH,
+            Action.INTERACT,
+            Direction.WEST,
+        ]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val3 = self.base_mdp.potential_function(state, mp)
-        
+
         # Pot onion
         if self.verbose:
             print("pot onion")
         actions = [Direction.WEST, Direction.NORTH, Action.INTERACT]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [Action.STAY, action])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [Action.STAY, action]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val4 = self.base_mdp.potential_function(state, mp)
 
         self.assertLess(val2, val3, "Potting tomato should increase potential")
@@ -491,7 +699,9 @@ def test_potential_function(self):
         # Pick up onion
         if self.verbose:
             print("pick up onion")
-        state, _ = self.base_mdp.get_state_transition(state, [Action.INTERACT, Action.STAY])
+        state, _ = self.base_mdp.get_state_transition(
+            state, [Action.INTERACT, Action.STAY]
+        )
         val5 = self.base_mdp.potential_function(state, mp)
         if self.verbose:
             print(self.base_mdp.state_string(state))
@@ -502,24 +712,42 @@ def test_potential_function(self):
             print("pick up tomato")
         actions = [Direction.SOUTH, Direction.EAST, Action.INTERACT]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [Action.STAY, action])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [Action.STAY, action]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val6 = self.base_mdp.potential_function(state, mp)
 
-        self.assertLess(val4, val5, "Picking up onion should increase potential")
-        self.assertLess(val5, val6, "Picking up tomato should increase potential")
+        self.assertLess(
+            val4, val5, "Picking up onion should increase potential"
+        )
+        self.assertLess(
+            val5, val6, "Picking up tomato should increase potential"
+        )
 
         # Pot onion
         if self.verbose:
             print("pot onion")
-        actions = [Direction.SOUTH, Direction.EAST, Direction.SOUTH, Action.INTERACT, Direction.WEST]
+        actions = [
+            Direction.SOUTH,
+            Direction.EAST,
+            Direction.SOUTH,
+            Action.INTERACT,
+            Direction.WEST,
+        ]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val7 = self.base_mdp.potential_function(state, mp)
 
         # Pot tomato
@@ -527,28 +755,34 @@ def test_potential_function(self):
             print("pot tomato")
         actions = [Direction.WEST, Direction.SOUTH, Action.INTERACT]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [Action.STAY, action])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [Action.STAY, action]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val8 = self.base_mdp.potential_function(state, mp)
 
-        
-
         self.assertLess(val6, val7, "Potting onion should increase potential")
         self.assertLess(val7, val8, "Potting tomato should increase potential")
 
         ## Useless pickups ##
-        
+
         # pickup tomato
         if self.verbose:
             print("pickup tomato")
         actions = [Action.INTERACT, Direction.NORTH]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val9 = self.base_mdp.potential_function(state, mp)
 
         # pickup tomato
@@ -556,41 +790,64 @@ def test_potential_function(self):
             print("pickup tomato")
         actions = [Direction.EAST, Action.INTERACT]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [Action.STAY, action])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [Action.STAY, action]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val10 = self.base_mdp.potential_function(state, mp)
 
-        self.assertLessEqual(val9, val8, "Extraneous pickup should not increase potential")
-        self.assertLessEqual(val10, val8, "Extraneous pickup should not increase potential")
+        self.assertLessEqual(
+            val9, val8, "Extraneous pickup should not increase potential"
+        )
+        self.assertLessEqual(
+            val10, val8, "Extraneous pickup should not increase potential"
+        )
 
         ## Catastrophic soup failure ##
-        
+
         # pot tomato
         if self.verbose:
             print("pot catastrophic tomato")
         actions = [Direction.WEST, Direction.SOUTH, Action.INTERACT]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [Action.STAY, action])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [Action.STAY, action]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val11 = self.base_mdp.potential_function(state, mp)
 
-        self.assertLess(val11, val10, "Catastrophic potting should decrease potential")
+        self.assertLess(
+            val11, val10, "Catastrophic potting should decrease potential"
+        )
 
         ## Bonus soup creation
 
         # pick up onion
         if self.verbose:
             print("pick up onion")
-        actions = [Direction.NORTH, Action.INTERACT, Direction.WEST, Action.INTERACT]
+        actions = [
+            Direction.NORTH,
+            Action.INTERACT,
+            Direction.WEST,
+            Action.INTERACT,
+        ]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val12 = self.base_mdp.potential_function(state, mp)
 
         # pot onion
@@ -598,10 +855,14 @@ def test_potential_function(self):
             print("pot onion")
         actions = [Direction.EAST, Direction.NORTH, Action.INTERACT]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val13 = self.base_mdp.potential_function(state, mp)
 
         # Cook soup
@@ -609,15 +870,25 @@ def test_potential_function(self):
             print("cook soup")
         actions = [Action.INTERACT, Direction.WEST]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val14 = self.base_mdp.potential_function(state, mp)
 
-        self.assertLess(val11, val12, "Useful onion pickup should increase potential")
-        self.assertLess(val12, val13, "Potting useful onion should increase potential")
-        self.assertLess(val13, val14, "Cooking optimal soup should increase potential")
+        self.assertLess(
+            val11, val12, "Useful onion pickup should increase potential"
+        )
+        self.assertLess(
+            val12, val13, "Potting useful onion should increase potential"
+        )
+        self.assertLess(
+            val13, val14, "Cooking optimal soup should increase potential"
+        )
 
         ## Soup pickup ##
 
@@ -626,10 +897,14 @@ def test_potential_function(self):
             print("pick up dish")
         actions = [Direction.WEST, Direction.SOUTH, Action.INTERACT]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [Action.STAY, action])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [Action.STAY, action]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val15 = self.base_mdp.potential_function(state, mp)
 
         # Move towards pot
@@ -637,24 +912,38 @@ def test_potential_function(self):
             print("move towards pot")
         actions = [Direction.EAST, Direction.NORTH]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [Action.STAY, action])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [Action.STAY, action]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val16 = self.base_mdp.potential_function(state, mp)
 
         # Pickup soup
         if self.verbose:
             print("pickup soup")
-        state, _ = self.base_mdp.get_state_transition(state, [Action.STAY, Action.INTERACT])
+        state, _ = self.base_mdp.get_state_transition(
+            state, [Action.STAY, Action.INTERACT]
+        )
         if self.verbose:
             print(self.base_mdp.state_string(state))
             print("potential: ", self.base_mdp.potential_function(state, mp))
         val17 = self.base_mdp.potential_function(state, mp)
 
-        self.assertLess(val14, val15, "Useful dish pickups should increase potential")
-        self.assertLess(val15, val16, "Moving towards soup with dish should increase potential")
-        self.assertLess(val16, val17, "Picking up soup should increase potential")
+        self.assertLess(
+            val14, val15, "Useful dish pickups should increase potential"
+        )
+        self.assertLess(
+            val15,
+            val16,
+            "Moving towards soup with dish should increase potential",
+        )
+        self.assertLess(
+            val16, val17, "Picking up soup should increase potential"
+        )
 
         ## Removing failed soup from pot
 
@@ -663,19 +952,27 @@ def test_potential_function(self):
             print("move torwards failed soup")
         actions = [Direction.SOUTH, Direction.EAST, Direction.SOUTH]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val18 = self.base_mdp.potential_function(state, mp)
 
         # Cook failed soup
         actions = [Action.INTERACT]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val19 = self.base_mdp.potential_function(state, mp)
 
         # Pickup dish
@@ -683,10 +980,14 @@ def test_potential_function(self):
             print("pickup dish")
         actions = [Direction.WEST, Direction.SOUTH, Action.INTERACT]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val20 = self.base_mdp.potential_function(state, mp)
 
         # Move towards soup
@@ -694,26 +995,46 @@ def test_potential_function(self):
             print("move towards soup")
         actions = [Direction.EAST, Direction.SOUTH]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val21 = self.base_mdp.potential_function(state, mp)
 
-        self.assertLess(val17, val18, "Moving towards failed soup should increase potential")
-        self.assertLess(val18, val19, "Cooking failed soup should increase potential")
-        self.assertLess(val19, val20, "Dish pickup for failed soup is still useful")
-        self.assertLess(val20, val21, "Moving towars pertinant pot with dish should increase potential")
+        self.assertLess(
+            val17,
+            val18,
+            "Moving towards failed soup should increase potential",
+        )
+        self.assertLess(
+            val18, val19, "Cooking failed soup should increase potential"
+        )
+        self.assertLess(
+            val19, val20, "Dish pickup for failed soup is still useful"
+        )
+        self.assertLess(
+            val20,
+            val21,
+            "Moving towars pertinant pot with dish should increase potential",
+        )
 
         ## Deliver failed soup ##
 
         # Pickup failed soup
         actions = [Action.INTERACT]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val22 = self.base_mdp.potential_function(state, mp)
 
         # Move towards serving area
@@ -721,24 +1042,42 @@ def test_potential_function(self):
             print("move towards servering area")
         actions = [Direction.EAST, Direction.SOUTH]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [action, Action.STAY])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [action, Action.STAY]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val23 = self.base_mdp.potential_function(state, mp)
 
         # Move away from serving area
         if self.verbose:
             print("move away from serving area")
-        state, _ = self.base_mdp.get_state_transition(state, [Direction.NORTH, Action.STAY])
+        state, _ = self.base_mdp.get_state_transition(
+            state, [Direction.NORTH, Action.STAY]
+        )
         if self.verbose:
             print(self.base_mdp.state_string(state))
             print("potential: ", self.base_mdp.potential_function(state, mp))
         val24 = self.base_mdp.potential_function(state, mp)
 
-        self.assertLess(val21, val22, "Picking up failed soup should increase potential")
-        self.assertAlmostEqual(val23, val22, delta=0.2, msg="Moving to serve failed soup doesn't change potential much")
-        self.assertAlmostEqual(val23, val24, delta=0.2, msg="Moving away from serving area with failed soup doesn't change much")
+        self.assertLess(
+            val21, val22, "Picking up failed soup should increase potential"
+        )
+        self.assertAlmostEqual(
+            val23,
+            val22,
+            delta=0.2,
+            msg="Moving to serve failed soup doesn't change potential much",
+        )
+        self.assertAlmostEqual(
+            val23,
+            val24,
+            delta=0.2,
+            msg="Moving away from serving area with failed soup doesn't change much",
+        )
 
         ## Deliver successful soup ##
 
@@ -747,25 +1086,36 @@ def test_potential_function(self):
             print("move towards serving area")
         actions = [Direction.SOUTH, Direction.EAST, Direction.SOUTH]
         for action in actions:
-            state, _ = self.base_mdp.get_state_transition(state, [Action.STAY, action])
+            state, _ = self.base_mdp.get_state_transition(
+                state, [Action.STAY, action]
+            )
             if self.verbose:
                 print(self.base_mdp.state_string(state))
-                print("potential: ", self.base_mdp.potential_function(state, mp))
+                print(
+                    "potential: ", self.base_mdp.potential_function(state, mp)
+                )
         val25 = self.base_mdp.potential_function(state, mp)
 
         # Deliver soup
         if self.verbose:
             print("deliver successful soup")
-        state, rewards = self.base_mdp.get_state_transition(state, [Action.STAY, Action.INTERACT])
+        state, rewards = self.base_mdp.get_state_transition(
+            state, [Action.STAY, Action.INTERACT]
+        )
         if self.verbose:
             print(self.base_mdp.state_string(state))
             print("potential: ", self.base_mdp.potential_function(state, mp))
 
-        self.assertLess(val24, val25, "Moving towards serving area with valid soup increases potential")
-        self.assertEqual(sum(rewards['sparse_reward_by_agent']), 50, "Soup was not properly devivered, probably an error with MDP logic")
-
-
-
+        self.assertLess(
+            val24,
+            val25,
+            "Moving towards serving area with valid soup increases potential",
+        )
+        self.assertEqual(
+            sum(rewards["sparse_reward_by_agent"]),
+            50,
+            "Soup was not properly devivered, probably an error with MDP logic",
+        )
 
 
 def random_joint_action():
@@ -775,99 +1125,167 @@ def random_joint_action():
 
 
 class TestFeaturizations(unittest.TestCase):
-
     def setUp(self):
         self.base_mdp = OvercookedGridworld.from_layout_name("cramped_room")
-        self.mlam = MediumLevelActionManager.from_pickle_or_compute(self.base_mdp, NO_COUNTERS_PARAMS, force_compute=True)
-        self.env = OvercookedEnv.from_mdp(self.base_mdp, **DEFAULT_ENV_PARAMS, info_level=0)
-        self.greedy_human_model_pair = AgentPair(GreedyHumanModel(self.mlam), GreedyHumanModel(self.mlam))
+        self.mlam = MediumLevelActionManager.from_pickle_or_compute(
+            self.base_mdp, NO_COUNTERS_PARAMS, force_compute=True
+        )
+        self.env = OvercookedEnv.from_mdp(
+            self.base_mdp, **DEFAULT_ENV_PARAMS, info_level=0
+        )
+        self.greedy_human_model_pair = AgentPair(
+            GreedyHumanModel(self.mlam), GreedyHumanModel(self.mlam)
+        )
         np.random.seed(0)
 
     def test_lossless_state_featurization_shape(self):
         s = self.base_mdp.get_standard_start_state()
         obs = self.base_mdp.lossless_state_encoding(s)[0]
-        self.assertTrue(np.array_equal(obs.shape, self.base_mdp.get_lossless_state_encoding_shape()), "{} vs {}".format(obs.shape, self.base_mdp.get_lossless_state_encoding_shape()))
+        self.assertTrue(
+            np.array_equal(
+                obs.shape, self.base_mdp.get_lossless_state_encoding_shape()
+            ),
+            "{} vs {}".format(
+                obs.shape, self.base_mdp.get_lossless_state_encoding_shape()
+            ),
+        )
 
     def test_state_featurization_shape(self):
         s = self.base_mdp.get_standard_start_state()
 
         for num_pots in range(3):
-            obs_0, obs_1 = self.base_mdp.featurize_state(s, self.mlam, num_pots=num_pots)
-            expected_shape = self.base_mdp.get_featurize_state_shape(num_pots=num_pots)
-            self.assertTrue(np.array_equal(obs_0.shape, expected_shape), "{} vs {}".format(obs_0.shape, expected_shape))
-            self.assertTrue(np.array_equal(obs_1.shape, expected_shape), "{} vs {}".format(obs_1.shape, expected_shape))
+            obs_0, obs_1 = self.base_mdp.featurize_state(
+                s, self.mlam, num_pots=num_pots
+            )
+            expected_shape = self.base_mdp.get_featurize_state_shape(
+                num_pots=num_pots
+            )
+            self.assertTrue(
+                np.array_equal(obs_0.shape, expected_shape),
+                "{} vs {}".format(obs_0.shape, expected_shape),
+            )
+            self.assertTrue(
+                np.array_equal(obs_1.shape, expected_shape),
+                "{} vs {}".format(obs_1.shape, expected_shape),
+            )
 
     def test_lossless_state_featurization(self):
-        trajs = self.env.get_rollouts(self.greedy_human_model_pair, num_games=5, info=False)
-        featurized_observations = [[self.base_mdp.lossless_state_encoding(state) for state in ep_states] for ep_states in trajs["ep_states"]]
-        
-        pickle_path = os.path.join(TESTING_DATA_DIR, "test_lossless_state_featurization", "expected")
+        trajs = self.env.get_rollouts(
+            self.greedy_human_model_pair, num_games=5, info=False
+        )
+        featurized_observations = [
+            [
+                self.base_mdp.lossless_state_encoding(state)
+                for state in ep_states
+            ]
+            for ep_states in trajs["ep_states"]
+        ]
+
+        pickle_path = os.path.join(
+            TESTING_DATA_DIR, "test_lossless_state_featurization", "expected"
+        )
 
         # NOTE: If the featurizations are updated intentionally, you can overwrite the expected
         # featurizations by uncommenting the following line:
         # save_pickle(featurized_observations, pickle_path)
         expected_featurization = load_pickle(pickle_path)
-        self.assertTrue(np.array_equal(expected_featurization, featurized_observations))
+        self.assertTrue(
+            np.array_equal(expected_featurization, featurized_observations)
+        )
 
     def test_state_featurization(self):
-        trajs = self.env.get_rollouts(self.greedy_human_model_pair, num_games=5, info=False)
+        trajs = self.env.get_rollouts(
+            self.greedy_human_model_pair, num_games=5, info=False
+        )
 
         for num_pots in range(3):
-            featurized_observations = [[self.base_mdp.featurize_state(state, self.mlam, num_pots=num_pots) for state in ep_states] for ep_states in trajs["ep_states"]]
-            pickle_path = os.path.join(TESTING_DATA_DIR, "test_state_featurization", 'expected_{}'.format(num_pots))
+            featurized_observations = [
+                [
+                    self.base_mdp.featurize_state(
+                        state, self.mlam, num_pots=num_pots
+                    )
+                    for state in ep_states
+                ]
+                for ep_states in trajs["ep_states"]
+            ]
+            pickle_path = os.path.join(
+                TESTING_DATA_DIR,
+                "test_state_featurization",
+                "expected_{}".format(num_pots),
+            )
             # NOTE: If the featurizations are updated intentionally, you can overwrite the expected
             # featurizations by uncommenting the following line:
             # save_pickle(featurized_observations, pickle_path)
             expected_featurization = load_pickle(pickle_path)
-            self.assertTrue(np.array_equal(expected_featurization, featurized_observations))
+            self.assertTrue(
+                np.array_equal(expected_featurization, featurized_observations)
+            )
 
     def test_state_featurization_symmetry(self):
-        trajs = self.env.get_rollouts(self.greedy_human_model_pair, num_games=5, info=False)
-        states = [state for ep_states in trajs['ep_states'] for state in ep_states]
-        
+        trajs = self.env.get_rollouts(
+            self.greedy_human_model_pair, num_games=5, info=False
+        )
+        states = [
+            state for ep_states in trajs["ep_states"] for state in ep_states
+        ]
+
         for state in states:
-            p0_obs, p1_obs = self.base_mdp.featurize_state(state, self.mlam, num_pots=2)
+            p0_obs, p1_obs = self.base_mdp.featurize_state(
+                state, self.mlam, num_pots=2
+            )
             state = state.reverse_players()
-            p0_obs_swapped, p1_obs_swapped = self.base_mdp.featurize_state(state, self.mlam, num_pots=2)
+            p0_obs_swapped, p1_obs_swapped = self.base_mdp.featurize_state(
+                state, self.mlam, num_pots=2
+            )
             state = state.reverse_players()
 
             self.assertTrue(np.array_equal(p0_obs, p1_obs_swapped))
             self.assertTrue(np.array_equal(p1_obs, p0_obs_swapped))
 
     def test_lossless_state_featurization_symmetry(self):
-        trajs = self.env.get_rollouts(self.greedy_human_model_pair, num_games=5, info=False)
-        states = [state for ep_states in trajs['ep_states'] for state in ep_states]
-        
+        trajs = self.env.get_rollouts(
+            self.greedy_human_model_pair, num_games=5, info=False
+        )
+        states = [
+            state for ep_states in trajs["ep_states"] for state in ep_states
+        ]
+
         for state in states:
             p0_obs, p1_obs = self.base_mdp.lossless_state_encoding(state)
             state = state.reverse_players()
-            p0_obs_swapped, p1_obs_swapped = self.base_mdp.lossless_state_encoding(state)
+            (
+                p0_obs_swapped,
+                p1_obs_swapped,
+            ) = self.base_mdp.lossless_state_encoding(state)
             state = state.reverse_players()
 
             self.assertTrue(np.array_equal(p0_obs, p1_obs_swapped))
             self.assertTrue(np.array_equal(p1_obs, p0_obs_swapped))
 
 
-
 class TestOvercookedEnvironment(unittest.TestCase):
 
-    dummy_dir = 'overcooked_test_temp'
+    dummy_dir = "overcooked_test_temp"
 
     def setUp(self):
         if not os.path.exists(self.dummy_dir):
             os.makedirs(self.dummy_dir)
         self.base_mdp = OvercookedGridworld.from_layout_name("cramped_room")
-        self.env = OvercookedEnv.from_mdp(self.base_mdp, info_level=0, **DEFAULT_ENV_PARAMS)
-        self.rnd_agent_pair = AgentPair(FixedPlanAgent([stay, w, w]), FixedPlanAgent([stay, e, e]))
+        self.env = OvercookedEnv.from_mdp(
+            self.base_mdp, info_level=0, **DEFAULT_ENV_PARAMS
+        )
+        self.rnd_agent_pair = AgentPair(
+            FixedPlanAgent([stay, w, w]), FixedPlanAgent([stay, e, e])
+        )
         np.random.seed(0)
 
     def tearDown(self):
         shutil.rmtree(self.dummy_dir)
 
     def _assert_files_equal(self, file_1, file_2):
-        with open(file_1, 'r') as f:
+        with open(file_1, "r") as f:
             lines_1 = f.readlines()
-        with open(file_2, 'r') as f:
+        with open(file_2, "r") as f:
             lines_2 = f.readlines()
 
         for line_1, line_2 in zip(lines_1, lines_2):
@@ -880,7 +1298,9 @@ def test_constructor(self):
             self.fail("Failed to instantiate OvercookedEnv:\n{}".format(e))
 
         with self.assertRaises(TypeError):
-            OvercookedEnv.from_mdp(self.base_mdp, **{"invalid_env_param": None})
+            OvercookedEnv.from_mdp(
+                self.base_mdp, **{"invalid_env_param": None}
+            )
 
     def test_step_fn(self):
         for _ in range(10):
@@ -889,7 +1309,9 @@ def test_step_fn(self):
 
     def test_execute_plan(self):
         action_plan = [random_joint_action() for _ in range(10)]
-        self.env.execute_plan(self.base_mdp.get_standard_start_state(), action_plan)
+        self.env.execute_plan(
+            self.base_mdp.get_standard_start_state(), action_plan
+        )
 
     def test_run_agents(self):
         start_state = self.env.state
@@ -909,34 +1331,68 @@ def test_one_player_env(self):
         a0 = FixedPlanAgent([stay, w, w, e, e, n, e, interact, w, n, interact])
         ag = AgentGroup(a0)
         env.run_agents(ag, display=False)
-        self.assertEqual(
-            env.state.players_pos_and_or,
-            (((2, 1), (0, -1)),)
-        )
+        self.assertEqual(env.state.players_pos_and_or, (((2, 1), (0, -1)),))
 
     def test_four_player_env_fixed(self):
         mdp = OvercookedGridworld.from_layout_name("multiplayer_schelling")
         assert mdp.num_players == 4
         env = OvercookedEnv.from_mdp(mdp, horizon=16, info_level=0)
         a0 = FixedPlanAgent([stay, w, w])
-        a1 = FixedPlanAgent([stay, stay, e, e, n, n, n, e, interact, n, n, w, w, w, n, interact, e])
-        a2 = FixedPlanAgent([stay, w, interact, n, n, e, e, e, n, e, n, interact, w])
+        a1 = FixedPlanAgent(
+            [
+                stay,
+                stay,
+                e,
+                e,
+                n,
+                n,
+                n,
+                e,
+                interact,
+                n,
+                n,
+                w,
+                w,
+                w,
+                n,
+                interact,
+                e,
+            ]
+        )
+        a2 = FixedPlanAgent(
+            [stay, w, interact, n, n, e, e, e, n, e, n, interact, w]
+        )
         a3 = FixedPlanAgent([e, interact, n, n, w, w, w, n, interact, e, s])
         ag = AgentGroup(a0, a1, a2, a3)
         env.run_agents(ag, display=False)
         self.assertEqual(
             env.state.players_pos_and_or,
-            (((1, 1), (-1, 0)), ((3, 1), (0, -1)), ((2, 1), (-1, 0)), ((4, 2), (0, 1)))
+            (
+                ((1, 1), (-1, 0)),
+                ((3, 1), (0, -1)),
+                ((2, 1), (-1, 0)),
+                ((4, 2), (0, 1)),
+            ),
         )
 
     def test_display(self):
         mdp0 = OvercookedGridworld.from_layout_name("cramped_room")
         mdp_fn = lambda _ignored: mdp0
         env = OvercookedEnv(mdp_fn, horizon=20)
-        env.get_rollouts(self.rnd_agent_pair, 1, display=True, info=False, dir=self.dummy_dir)
+        env.get_rollouts(
+            self.rnd_agent_pair,
+            1,
+            display=True,
+            info=False,
+            dir=self.dummy_dir,
+        )
 
-        expected_display_file = os.path.join(TESTING_DATA_DIR, 'test_display', 'expected.txt')
-        actual_display_file = glob.glob(os.path.join(self.dummy_dir, '*.txt'))[0]
+        expected_display_file = os.path.join(
+            TESTING_DATA_DIR, "test_display", "expected.txt"
+        )
+        actual_display_file = glob.glob(os.path.join(self.dummy_dir, "*.txt"))[
+            0
+        ]
 
         # If display intentionally updated, uncomment the line below to update expected values
         shutil.copy(actual_display_file, expected_display_file)
@@ -946,10 +1402,21 @@ def test_display_phi(self):
         mdp0 = OvercookedGridworld.from_layout_name("cramped_room")
         mdp_fn = lambda _ignored: mdp0
         env = OvercookedEnv(mdp_fn, horizon=20)
-        env.get_rollouts(self.rnd_agent_pair, 1, display=True, display_phi=True, info=False, dir=self.dummy_dir)
+        env.get_rollouts(
+            self.rnd_agent_pair,
+            1,
+            display=True,
+            display_phi=True,
+            info=False,
+            dir=self.dummy_dir,
+        )
 
-        expected_display_file = os.path.join(TESTING_DATA_DIR, 'test_display_phi', 'expected.txt')
-        actual_display_file = glob.glob(os.path.join(self.dummy_dir, '*.txt'))[0]
+        expected_display_file = os.path.join(
+            TESTING_DATA_DIR, "test_display_phi", "expected.txt"
+        )
+        actual_display_file = glob.glob(os.path.join(self.dummy_dir, "*.txt"))[
+            0
+        ]
 
         # If display intentionally updated, uncomment the line below to update expected values
         shutil.copy(actual_display_file, expected_display_file)
@@ -959,15 +1426,19 @@ def test_multiple_mdp_env(self):
         mdp0 = OvercookedGridworld.from_layout_name("cramped_room")
         mdp1 = OvercookedGridworld.from_layout_name("counter_circuit")
         mdp_fn = lambda _ignored: np.random.choice([mdp0, mdp1])
-        
+
         # Default env
         env = OvercookedEnv(mdp_fn, horizon=100)
         env.get_rollouts(self.rnd_agent_pair, 5, info=False)
 
     def test_starting_position_randomization(self):
         self.base_mdp = OvercookedGridworld.from_layout_name("cramped_room")
-        start_state_fn = self.base_mdp.get_random_start_state_fn(random_start_pos=True, rnd_obj_prob_thresh=0.0)
-        env = OvercookedEnv.from_mdp(self.base_mdp, start_state_fn, info_level=0)
+        start_state_fn = self.base_mdp.get_random_start_state_fn(
+            random_start_pos=True, rnd_obj_prob_thresh=0.0
+        )
+        env = OvercookedEnv.from_mdp(
+            self.base_mdp, start_state_fn, info_level=0
+        )
         start_state = env.state.players_pos_and_or
         for _ in range(3):
             env.reset()
@@ -976,8 +1447,12 @@ def test_starting_position_randomization(self):
 
     def test_starting_obj_randomization(self):
         self.base_mdp = OvercookedGridworld.from_layout_name("cramped_room")
-        start_state_fn = self.base_mdp.get_random_start_state_fn(random_start_pos=False, rnd_obj_prob_thresh=0.8)
-        env = OvercookedEnv.from_mdp(self.base_mdp, start_state_fn, info_level=0)
+        start_state_fn = self.base_mdp.get_random_start_state_fn(
+            random_start_pos=False, rnd_obj_prob_thresh=0.8
+        )
+        env = OvercookedEnv.from_mdp(
+            self.base_mdp, start_state_fn, info_level=0
+        )
         start_state = env.state.all_objects_list
         for _ in range(3):
             env.reset()
@@ -991,17 +1466,18 @@ def test_failing_rnd_layout(self):
             OvercookedEnv(mdp_fn, **DEFAULT_ENV_PARAMS)
 
     def test_random_layout(self):
-        mdp_gen_params = {"inner_shape": (5, 4),
-                            "prop_empty": 0.8,
-                            "prop_feats": 0.2,
-                            "start_all_orders" : [
-                                { "ingredients" : ["onion", "onion", "onion"]}
-                            ],
-                            "recipe_values" : [20],
-                            "recipe_times" : [20],
-                            "display": False
-                          }
-        mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(mdp_gen_params, outer_shape=(5, 4))
+        mdp_gen_params = {
+            "inner_shape": (5, 4),
+            "prop_empty": 0.8,
+            "prop_feats": 0.2,
+            "start_all_orders": [{"ingredients": ["onion", "onion", "onion"]}],
+            "recipe_values": [20],
+            "recipe_times": [20],
+            "display": False,
+        }
+        mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(
+            mdp_gen_params, outer_shape=(5, 4)
+        )
         env = OvercookedEnv(mdp_fn, **DEFAULT_ENV_PARAMS)
         start_terrain = env.mdp.terrain_mtx
 
@@ -1010,106 +1486,172 @@ def test_random_layout(self):
             curr_terrain = env.mdp.terrain_mtx
             self.assertFalse(np.array_equal(start_terrain, curr_terrain))
 
-        mdp_gen_params = {"layout_name": 'cramped_room'}
+        mdp_gen_params = {"layout_name": "cramped_room"}
         mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(mdp_gen_params)
         env = OvercookedEnv(mdp_fn, **DEFAULT_ENV_PARAMS)
-        
+
         layouts_seen = []
         for _ in range(5):
             layouts_seen.append(env.mdp.terrain_mtx)
             env.reset()
 
-        all_same_layout = all([np.array_equal(env.mdp.terrain_mtx, terrain) for terrain in layouts_seen])
+        all_same_layout = all(
+            [
+                np.array_equal(env.mdp.terrain_mtx, terrain)
+                for terrain in layouts_seen
+            ]
+        )
         self.assertTrue(all_same_layout)
 
-        mdp_gen_params = {"layout_name": 'asymmetric_advantages'}
+        mdp_gen_params = {"layout_name": "asymmetric_advantages"}
         mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(mdp_gen_params)
         env = OvercookedEnv(mdp_fn, **DEFAULT_ENV_PARAMS)
         for _ in range(5):
             layouts_seen.append(env.mdp.terrain_mtx)
             env.reset()
 
-        all_same_layout = all([np.array_equal(env.mdp.terrain_mtx, terrain) for terrain in layouts_seen])
+        all_same_layout = all(
+            [
+                np.array_equal(env.mdp.terrain_mtx, terrain)
+                for terrain in layouts_seen
+            ]
+        )
         self.assertFalse(all_same_layout)
-        
+
     def test_random_layout_feature_types(self):
         mandatory_features = {POT, DISH_DISPENSER, SERVING_LOC}
         optional_features = {ONION_DISPENSER, TOMATO_DISPENSER}
-        optional_features_combinations = [{ONION_DISPENSER, TOMATO_DISPENSER}, {ONION_DISPENSER}, {TOMATO_DISPENSER}]
+        optional_features_combinations = [
+            {ONION_DISPENSER, TOMATO_DISPENSER},
+            {ONION_DISPENSER},
+            {TOMATO_DISPENSER},
+        ]
 
         for optional_features_combo in optional_features_combinations:
-            left_out_optional_features = optional_features - optional_features_combo
+            left_out_optional_features = (
+                optional_features - optional_features_combo
+            )
             used_features = list(optional_features_combo | mandatory_features)
-            mdp_gen_params = {"prop_feats": 0.9,
-                            "feature_types": used_features,
-                            "prop_empty": 0.1,
-                            "inner_shape": (6, 5),
-                            "display": False,
-                            "start_all_orders" : [
-                                { "ingredients" : ["onion", "onion", "onion"]}
-                            ]}
-            mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(mdp_gen_params, outer_shape=(6, 5))
+            mdp_gen_params = {
+                "prop_feats": 0.9,
+                "feature_types": used_features,
+                "prop_empty": 0.1,
+                "inner_shape": (6, 5),
+                "display": False,
+                "start_all_orders": [
+                    {"ingredients": ["onion", "onion", "onion"]}
+                ],
+            }
+            mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(
+                mdp_gen_params, outer_shape=(6, 5)
+            )
             env = OvercookedEnv(mdp_fn, **DEFAULT_ENV_PARAMS)
             for _ in range(10):
                 env.reset()
                 curr_terrain = env.mdp.terrain_mtx
-                terrain_features = set.union(*(set(line) for line in curr_terrain))
-                self.assertTrue(all(elem in terrain_features for elem in used_features)) # all used_features are actually used
+                terrain_features = set.union(
+                    *(set(line) for line in curr_terrain)
+                )
+                self.assertTrue(
+                    all(elem in terrain_features for elem in used_features)
+                )  # all used_features are actually used
                 if left_out_optional_features:
-                    self.assertFalse(any(elem in terrain_features for elem in left_out_optional_features)) # all left_out optional_features are not used
+                    self.assertFalse(
+                        any(
+                            elem in terrain_features
+                            for elem in left_out_optional_features
+                        )
+                    )  # all left_out optional_features are not used
 
     def test_random_layout_generated_recipes(self):
-        only_onions_recipes = [Recipe(["onion", "onion"]), Recipe(["onion", "onion", "onion"])]
+        only_onions_recipes = [
+            Recipe(["onion", "onion"]),
+            Recipe(["onion", "onion", "onion"]),
+        ]
         only_onions_dict_recipes = [r.to_dict() for r in only_onions_recipes]
 
         # checking if recipes are generated from mdp_params
-        mdp_gen_params = {"generate_all_orders": {"n":2, "ingredients": ["onion"], "min_size":2, "max_size":3},
-                        "prop_feats": 0.9,
-                        "prop_empty": 0.1,
-                        "inner_shape": (6, 5),
-                        "display": False}
-        mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(mdp_gen_params, outer_shape=(6, 5))
+        mdp_gen_params = {
+            "generate_all_orders": {
+                "n": 2,
+                "ingredients": ["onion"],
+                "min_size": 2,
+                "max_size": 3,
+            },
+            "prop_feats": 0.9,
+            "prop_empty": 0.1,
+            "inner_shape": (6, 5),
+            "display": False,
+        }
+        mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(
+            mdp_gen_params, outer_shape=(6, 5)
+        )
         env = OvercookedEnv(mdp_fn, **DEFAULT_ENV_PARAMS)
         for _ in range(10):
             env.reset()
-            self.assertCountEqual(env.mdp.start_all_orders, only_onions_dict_recipes)
+            self.assertCountEqual(
+                env.mdp.start_all_orders, only_onions_dict_recipes
+            )
             self.assertEqual(len(env.mdp.start_bonus_orders), 0)
-        
+
         # checking if bonus_orders is subset of all_orders even if not specified
 
-        mdp_gen_params = {"generate_all_orders": {"n":2, "ingredients": ["onion"], "min_size":2, "max_size":3},
-                        "generate_bonus_orders": {"n":1, "min_size":2, "max_size":3},
-                        "prop_feats": 0.9,
-                        "prop_empty": 0.1,
-                        "inner_shape": (6, 5),
-                        "display": False}
-        mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(mdp_gen_params, outer_shape=(6,5))
+        mdp_gen_params = {
+            "generate_all_orders": {
+                "n": 2,
+                "ingredients": ["onion"],
+                "min_size": 2,
+                "max_size": 3,
+            },
+            "generate_bonus_orders": {"n": 1, "min_size": 2, "max_size": 3},
+            "prop_feats": 0.9,
+            "prop_empty": 0.1,
+            "inner_shape": (6, 5),
+            "display": False,
+        }
+        mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(
+            mdp_gen_params, outer_shape=(6, 5)
+        )
         env = OvercookedEnv(mdp_fn, **DEFAULT_ENV_PARAMS)
         for _ in range(10):
             env.reset()
-            self.assertCountEqual(env.mdp.start_all_orders, only_onions_dict_recipes)
+            self.assertCountEqual(
+                env.mdp.start_all_orders, only_onions_dict_recipes
+            )
             self.assertEqual(len(env.mdp.start_bonus_orders), 1)
-            self.assertTrue(env.mdp.start_bonus_orders[0] in only_onions_dict_recipes)
+            self.assertTrue(
+                env.mdp.start_bonus_orders[0] in only_onions_dict_recipes
+            )
 
         # checking if after reset there are new recipes generated
-        mdp_gen_params = {"generate_all_orders": {"n":3, "min_size":2, "max_size":3},
-                        "prop_feats": 0.9,
-                        "prop_empty": 0.1,
-                        "inner_shape": (6, 5),
-                        "display": False,
-                        "feature_types": [POT, DISH_DISPENSER, SERVING_LOC, ONION_DISPENSER, TOMATO_DISPENSER]
-                        }
-        mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(mdp_gen_params, outer_shape=(6,5))
+        mdp_gen_params = {
+            "generate_all_orders": {"n": 3, "min_size": 2, "max_size": 3},
+            "prop_feats": 0.9,
+            "prop_empty": 0.1,
+            "inner_shape": (6, 5),
+            "display": False,
+            "feature_types": [
+                POT,
+                DISH_DISPENSER,
+                SERVING_LOC,
+                ONION_DISPENSER,
+                TOMATO_DISPENSER,
+            ],
+        }
+        mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(
+            mdp_gen_params, outer_shape=(6, 5)
+        )
         env = OvercookedEnv(mdp_fn, **DEFAULT_ENV_PARAMS)
         generated_recipes_strings = set()
         for _ in range(20):
             env.reset()
-            generated_recipes_strings |= {json.dumps(o, sort_keys=True) for o in env.mdp.start_all_orders}
+            generated_recipes_strings |= {
+                json.dumps(o, sort_keys=True) for o in env.mdp.start_all_orders
+            }
         self.assertTrue(len(generated_recipes_strings) > 3)
-        
-        
-class TestGymEnvironment(unittest.TestCase):
 
+
+class TestGymEnvironment(unittest.TestCase):
     def setUp(self):
         self.base_mdp = OvercookedGridworld.from_layout_name("cramped_room")
         self.env = OvercookedEnv.from_mdp(self.base_mdp, **DEFAULT_ENV_PARAMS)
@@ -1118,18 +1660,28 @@ def setUp(self):
 
     # TODO: write more tests here
 
-class TestTrajectories(unittest.TestCase):
 
+class TestTrajectories(unittest.TestCase):
     def setUp(self):
         self.base_mdp = OvercookedGridworld.from_layout_name("cramped_room")
-        self.mlam = MediumLevelActionManager.from_pickle_or_compute(self.base_mdp, NO_COUNTERS_PARAMS, force_compute=True)
-        self.env = OvercookedEnv.from_mdp(self.base_mdp, **DEFAULT_ENV_PARAMS, info_level=0)
-        self.greedy_human_model_pair = AgentPair(GreedyHumanModel(self.mlam), GreedyHumanModel(self.mlam))
+        self.mlam = MediumLevelActionManager.from_pickle_or_compute(
+            self.base_mdp, NO_COUNTERS_PARAMS, force_compute=True
+        )
+        self.env = OvercookedEnv.from_mdp(
+            self.base_mdp, **DEFAULT_ENV_PARAMS, info_level=0
+        )
+        self.greedy_human_model_pair = AgentPair(
+            GreedyHumanModel(self.mlam), GreedyHumanModel(self.mlam)
+        )
         np.random.seed(0)
 
     def test_append(self):
-        traj_one = self.env.get_rollouts(self.greedy_human_model_pair, num_games=3, info=False)
-        traj_two = self.env.get_rollouts(self.greedy_human_model_pair, num_games=3, info=False)
+        traj_one = self.env.get_rollouts(
+            self.greedy_human_model_pair, num_games=3, info=False
+        )
+        traj_two = self.env.get_rollouts(
+            self.greedy_human_model_pair, num_games=3, info=False
+        )
 
         combined = append_trajectories(traj_one, traj_two)
 
@@ -1150,5 +1702,6 @@ def test_append(self):
             shape_combined = combined[key].shape
             self.assertEqual(shape_combined[0], shape_one[0] + shape_two[0])
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     unittest.main()
diff --git a/testing/planners_test.py b/testing/planners_test.py
index 6e1c7d89..688b25fc 100644
--- a/testing/planners_test.py
+++ b/testing/planners_test.py
@@ -1,10 +1,17 @@
 import unittest
-from overcooked_ai_py.planning.planners import MediumLevelActionManager
-from overcooked_ai_py.mdp.actions import Direction, Action
-from overcooked_ai_py.mdp.overcooked_mdp import OvercookedGridworld, PlayerState, ObjectState, SoupState, OvercookedState
-from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
-from overcooked_ai_py.agents.benchmarking import AgentEvaluator
+
 from overcooked_ai_py.agents.agent import AgentPair, GreedyHumanModel
+from overcooked_ai_py.agents.benchmarking import AgentEvaluator
+from overcooked_ai_py.mdp.actions import Action, Direction
+from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
+from overcooked_ai_py.mdp.overcooked_mdp import (
+    ObjectState,
+    OvercookedGridworld,
+    OvercookedState,
+    PlayerState,
+    SoupState,
+)
+from overcooked_ai_py.planning.planners import MediumLevelActionManager
 
 large_mdp_tests = False
 force_compute = True
@@ -16,63 +23,81 @@
 P, Obj = PlayerState, ObjectState
 
 # Simple MDP Setup
-simple_mdp = OvercookedGridworld.from_layout_name('simple_o')
+simple_mdp = OvercookedGridworld.from_layout_name("simple_o")
 
 
 base_params = {
-    'start_orientations': False,
-    'wait_allowed': False,
-    'counter_goals': simple_mdp.terrain_pos_dict['X'],
-    'counter_drop': simple_mdp.terrain_pos_dict['X'][1:2],
-    'counter_pickup': simple_mdp.terrain_pos_dict['X'],
-    'same_motion_goals': True
+    "start_orientations": False,
+    "wait_allowed": False,
+    "counter_goals": simple_mdp.terrain_pos_dict["X"],
+    "counter_drop": simple_mdp.terrain_pos_dict["X"][1:2],
+    "counter_pickup": simple_mdp.terrain_pos_dict["X"],
+    "same_motion_goals": True,
 }
 action_manger_filename = "simple_1_am.pkl"
 ml_action_manager_simple = MediumLevelActionManager.from_pickle_or_compute(
-    simple_mdp, mlam_params=base_params, custom_filename=action_manger_filename, force_compute=force_compute)
+    simple_mdp,
+    mlam_params=base_params,
+    custom_filename=action_manger_filename,
+    force_compute=force_compute,
+)
 ml_action_manager_simple.env = OvercookedEnv.from_mdp(simple_mdp)
 
 base_params_start_or = {
-    'start_orientations': True,
-    'wait_allowed': False,
-    'counter_goals': simple_mdp.terrain_pos_dict['X'],
-    'counter_drop': [],
-    'counter_pickup': simple_mdp.terrain_pos_dict['X'],
-    'same_motion_goals': False
+    "start_orientations": True,
+    "wait_allowed": False,
+    "counter_goals": simple_mdp.terrain_pos_dict["X"],
+    "counter_drop": [],
+    "counter_pickup": simple_mdp.terrain_pos_dict["X"],
+    "same_motion_goals": False,
 }
 action_manger_filename = "simple_2_am.pkl"
 or_ml_action_manager_simple = MediumLevelActionManager.from_pickle_or_compute(
-    simple_mdp, mlam_params=base_params_start_or, custom_filename=action_manger_filename, force_compute=force_compute)
+    simple_mdp,
+    mlam_params=base_params_start_or,
+    custom_filename=action_manger_filename,
+    force_compute=force_compute,
+)
 
 if large_mdp_tests:
     # Not testing by default
 
     # Large MDP Setup
-    large_mdp = OvercookedGridworld.from_layout_name('corridor', cook_time=5)
+    large_mdp = OvercookedGridworld.from_layout_name("corridor", cook_time=5)
 
     no_counters_params = {
-        'start_orientations': False,
-        'wait_allowed': False,
-        'counter_goals': [],
-        'counter_drop': [],
-        'counter_pickup': [],
-        'same_motion_goals': False
+        "start_orientations": False,
+        "wait_allowed": False,
+        "counter_goals": [],
+        "counter_drop": [],
+        "counter_pickup": [],
+        "same_motion_goals": False,
     }
     action_manger_filename = "corridor_no_shared_motion_goals_am.pkl"
-    ml_planner_large_no_shared = MediumLevelActionManager.from_pickle_or_compute(
-        large_mdp, no_counters_params, custom_filename=action_manger_filename, force_compute=force_compute_large)
+    ml_planner_large_no_shared = (
+        MediumLevelActionManager.from_pickle_or_compute(
+            large_mdp,
+            no_counters_params,
+            custom_filename=action_manger_filename,
+            force_compute=force_compute_large,
+        )
+    )
 
     same_goals_params = {
-        'start_orientations': False,
-        'wait_allowed': False,
-        'counter_goals': [],
-        'counter_drop': [],
-        'counter_pickup': [],
-        'same_motion_goals': True
+        "start_orientations": False,
+        "wait_allowed": False,
+        "counter_goals": [],
+        "counter_drop": [],
+        "counter_pickup": [],
+        "same_motion_goals": True,
     }
     action_manger_filename = "corridor_am.pkl"
     ml_planner_large = MediumLevelActionManager.from_pickle_or_compute(
-        large_mdp, same_goals_params, custom_filename=action_manger_filename, force_compute=force_compute_large)
+        large_mdp,
+        same_goals_params,
+        custom_filename=action_manger_filename,
+        force_compute=force_compute_large,
+    )
     # Deprecated.
     # hlam = HighLevelActionManager(ml_planner_large)
     # hlp = HighLevelPlanner(hlam)
@@ -93,12 +118,11 @@ def cooking_soup_obj(soup_loc, num_onion_inside=3, cooking_tick=0):
 
 
 def soup_obj(soup_loc, num_onion_inside, cooking_tick):
-    ingredient_obj_lst = [Obj('onion', soup_loc)] * num_onion_inside
+    ingredient_obj_lst = [Obj("onion", soup_loc)] * num_onion_inside
     return SoupState(soup_loc, ingredient_obj_lst, cooking_tick)
 
 
 class TestMotionPlanner(unittest.TestCase):
-
     def test_gridworld_distance(self):
         planner = ml_action_manager_simple.joint_motion_planner.motion_planner
         start = ((2, 1), e)
@@ -127,32 +151,44 @@ def test_simple_mdp(self):
         self.simple_mdp_orientation_change(planner)
         self.simple_mdp_basic_plan(planner)
         self.simple_mdp_orientation_optimization_dependent_plans(planner)
-        
+
     def simple_mdp_already_at_goal(self, planner):
         start_status = goal_status = ((1, 1), n)
-        self.check_single_motion_plan(planner, start_status, goal_status, expected_length=1)
+        self.check_single_motion_plan(
+            planner, start_status, goal_status, expected_length=1
+        )
 
     def simple_mdp_orientation_change(self, planner):
         start_status = ((1, 1), n)
         goal_status = ((1, 1), w)
-        self.check_single_motion_plan(planner, start_status, goal_status, expected_length=2)
+        self.check_single_motion_plan(
+            planner, start_status, goal_status, expected_length=2
+        )
 
     def simple_mdp_basic_plan(self, planner):
         start_status = ((1, 1), n)
         goal_status = ((3, 1), n)
-        self.check_single_motion_plan(planner, start_status, goal_status, expected_length=4)
+        self.check_single_motion_plan(
+            planner, start_status, goal_status, expected_length=4
+        )
 
     def simple_mdp_orientation_optimization_dependent_plans(self, planner):
         start_status = ((2, 1), n)
         goal_status = ((1, 2), w)
-        self.check_single_motion_plan(planner, start_status, goal_status, expected_length=3)
+        self.check_single_motion_plan(
+            planner, start_status, goal_status, expected_length=3
+        )
 
         goal_status = ((1, 2), s)
-        self.check_single_motion_plan(planner, start_status, goal_status, expected_length=3)
+        self.check_single_motion_plan(
+            planner, start_status, goal_status, expected_length=3
+        )
 
     def test_larger_mdp(self):
         if large_mdp_tests:
-            planner = ml_planner_large.ml_action_manager.joint_motion_planner.motion_planner
+            planner = (
+                ml_planner_large.ml_action_manager.joint_motion_planner.motion_planner
+            )
             self.large_mdp_basic_plan(planner)
 
     def large_mdp_basic_plan(self, planner):
@@ -160,39 +196,58 @@ def large_mdp_basic_plan(self, planner):
         goal_status = ((8, 1), n)
         self.check_single_motion_plan(planner, start_status, goal_status)
 
-    def check_single_motion_plan(self, motion_planner, start_pos_and_or, goal_pos_and_or, expected_length=None):
+    def check_single_motion_plan(
+        self,
+        motion_planner,
+        start_pos_and_or,
+        goal_pos_and_or,
+        expected_length=None,
+    ):
         dummy_agent = P((3, 2), n)
-        start_state = OvercookedState([P(*start_pos_and_or), dummy_agent], {}, all_orders=simple_mdp.start_all_orders)
-        action_plan, pos_and_or_plan, plan_cost = motion_planner.get_plan(start_pos_and_or, goal_pos_and_or)
-        
+        start_state = OvercookedState(
+            [P(*start_pos_and_or), dummy_agent],
+            {},
+            all_orders=simple_mdp.start_all_orders,
+        )
+        action_plan, pos_and_or_plan, plan_cost = motion_planner.get_plan(
+            start_pos_and_or, goal_pos_and_or
+        )
+
         # Checking that last state obtained matches goal position
         self.assertEqual(pos_and_or_plan[-1], goal_pos_and_or)
 
         # In single motion plans the graph cost should be equal to
         # the plan cost (= plan length) as agents should never STAY
-        graph_plan_cost = sum([motion_planner._graph_action_cost(a) for a in action_plan])
+        graph_plan_cost = sum(
+            [motion_planner._graph_action_cost(a) for a in action_plan]
+        )
         self.assertEqual(plan_cost, graph_plan_cost)
 
         joint_action_plan = [(a, stay) for a in action_plan]
         env = OvercookedEnv.from_mdp(motion_planner.mdp, horizon=1000)
         resulting_state, _ = env.execute_plan(start_state, joint_action_plan)
-        self.assertEqual(resulting_state.players_pos_and_or[0], goal_pos_and_or)
+        self.assertEqual(
+            resulting_state.players_pos_and_or[0], goal_pos_and_or
+        )
 
-        if expected_length is not None: 
+        if expected_length is not None:
             self.assertEqual(len(action_plan), expected_length)
 
 
 class TestJointMotionPlanner(unittest.TestCase):
-
     def test_same_start_and_end_pos_with_no_start_orientations(self):
         jm_planner = ml_action_manager_simple.joint_motion_planner
         start = (((1, 1), w), ((1, 2), s))
         goal = (((1, 1), n), ((2, 1), n))
 
-        joint_action_plan, end_jm_state, finshing_times = jm_planner.get_low_level_action_plan(start, goal)
+        (
+            joint_action_plan,
+            end_jm_state,
+            finshing_times,
+        ) = jm_planner.get_low_level_action_plan(start, goal)
         optimal_plan = [(n, e), (interact, n)]
         self.assertEqual(joint_action_plan, optimal_plan)
-        
+
         optimal_end_jm_state = (((1, 1), n), ((2, 1), n))
         self.assertEqual(end_jm_state, optimal_end_jm_state)
 
@@ -202,7 +257,7 @@ def test_same_start_and_end_pos_with_no_start_orientations(self):
     def test_with_start_orientations_simple_mdp(self):
         jm_planner = or_ml_action_manager_simple.joint_motion_planner
         self.simple_mdp_suite(jm_planner)
-    
+
     def test_without_start_orientations_simple_mdp(self):
         jm_planner = ml_action_manager_simple.joint_motion_planner
         self.simple_mdp_suite(jm_planner)
@@ -325,22 +380,56 @@ def large_mdp_test_shared_motion_goal_with_conflict_other(self, planner):
         goal = (a1_goal, a2_goal)
         self.check_joint_plan(planner, start, goal, min_t=3)
 
-    def check_joint_plan(self, joint_motion_planner, start, goal, times=None, min_t=None, display=False):
+    def check_joint_plan(
+        self,
+        joint_motion_planner,
+        start,
+        goal,
+        times=None,
+        min_t=None,
+        display=False,
+    ):
         """Runs the plan in the environment and checks that the intended goals are achieved."""
         debug = False
-        action_plan, end_pos_and_orients, plan_lengths = joint_motion_planner.get_low_level_action_plan(start, goal)
-        if debug: print("Start state: {}, Goal state: {}, Action plan: {}".format(start, goal, action_plan))
-
-        start_state = OvercookedState([P(*start[0]), P(*start[1])], {}, all_orders=simple_mdp.start_all_orders)
+        (
+            action_plan,
+            end_pos_and_orients,
+            plan_lengths,
+        ) = joint_motion_planner.get_low_level_action_plan(start, goal)
+        if debug:
+            print(
+                "Start state: {}, Goal state: {}, Action plan: {}".format(
+                    start, goal, action_plan
+                )
+            )
+
+        start_state = OvercookedState(
+            [P(*start[0]), P(*start[1])],
+            {},
+            all_orders=simple_mdp.start_all_orders,
+        )
         env = OvercookedEnv.from_mdp(joint_motion_planner.mdp, horizon=1000)
-        resulting_state, _ = env.execute_plan(start_state, action_plan, display=display)
-
-        self.assertTrue(any([agent_goal in resulting_state.players_pos_and_or for agent_goal in goal]))
-        self.assertEqual(resulting_state.players_pos_and_or, end_pos_and_orients)
+        resulting_state, _ = env.execute_plan(
+            start_state, action_plan, display=display
+        )
+
+        self.assertTrue(
+            any(
+                [
+                    agent_goal in resulting_state.players_pos_and_or
+                    for agent_goal in goal
+                ]
+            )
+        )
+        self.assertEqual(
+            resulting_state.players_pos_and_or, end_pos_and_orients
+        )
         self.assertEqual(len(action_plan), min(plan_lengths))
 
-        if min_t is not None: self.assertEqual(len(action_plan), min_t)
-        if times is not None: self.assertEqual(plan_lengths, times)
+        if min_t is not None:
+            self.assertEqual(len(action_plan), min_t)
+        if times is not None:
+            self.assertEqual(plan_lengths, times)
 
 
 # Rewritten because the previous test depended on Heuristic, and Heuristic has been deprecated
@@ -373,14 +462,26 @@ def test_simple_mdp_with_start_orientations(self):
         self.simple_mdp_pickup_counter_soup(mlam, counter_drop_forbidden=True)
         self.simple_mdp_pickup_counter_dish(mlam, counter_drop_forbidden=True)
         self.simple_mdp_pickup_counter_onion(mlam, counter_drop_forbidden=True)
-        self.simple_mdp_drop_useless_dish_with_soup_idle(mlam, counter_drop_forbidden=True)
+        self.simple_mdp_drop_useless_dish_with_soup_idle(
+            mlam, counter_drop_forbidden=True
+        )
         self.simple_mdp_pickup_soup(mlam, counter_drop_forbidden=True)
         self.simple_mdp_pickup_dish(mlam, counter_drop_forbidden=True)
-        self.simple_mdp_start_good_soup_cooking(mlam, counter_drop_forbidden=True)
-        self.simple_mdp_start_bad_soup_cooking(mlam, counter_drop_forbidden=True)
-        self.simple_mdp_start_1_onion_soup_cooking(mlam, counter_drop_forbidden=True)
-        self.simple_mdp_drop_useless_onion_good_soup(mlam, counter_drop_forbidden=True)
-        self.simple_mdp_drop_useless_onion_bad_soup(mlam, counter_drop_forbidden=True)
+        self.simple_mdp_start_good_soup_cooking(
+            mlam, counter_drop_forbidden=True
+        )
+        self.simple_mdp_start_bad_soup_cooking(
+            mlam, counter_drop_forbidden=True
+        )
+        self.simple_mdp_start_1_onion_soup_cooking(
+            mlam, counter_drop_forbidden=True
+        )
+        self.simple_mdp_drop_useless_onion_good_soup(
+            mlam, counter_drop_forbidden=True
+        )
+        self.simple_mdp_drop_useless_onion_bad_soup(
+            mlam, counter_drop_forbidden=True
+        )
         self.simple_mdp_add_3rd_onion(mlam, counter_drop_forbidden=True)
         self.simple_mdp_add_2nd_onion(mlam, counter_drop_forbidden=True)
         self.simple_mdp_drop_useless_dish(mlam, counter_drop_forbidden=True)
@@ -394,230 +495,294 @@ def test_simple_mdp_with_start_orientations(self):
 
     def simple_mpd_empty_hands(self, planner, counter_drop_forbidden=False):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n)],
+            [P((2, 2), n), P((2, 1), n)],
             {},
-            all_orders=simple_mdp.start_all_orders)
-        self.check_ml_action_manager(s, planner,
-                                     [self.ONION_PICKUP, self.DISH_PICKUP],
-                                     [self.ONION_PICKUP, self.DISH_PICKUP]
-                                     )
+            all_orders=simple_mdp.start_all_orders,
+        )
+        self.check_ml_action_manager(
+            s,
+            planner,
+            [self.ONION_PICKUP, self.DISH_PICKUP],
+            [self.ONION_PICKUP, self.DISH_PICKUP],
+        )
 
     def simple_mdp_deliver_soup(self, planner, counter_drop_forbidden=False):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n, done_soup_obj((2, 1)))],
+            [P((2, 2), n), P((2, 1), n, done_soup_obj((2, 1)))],
             {},
-            all_orders=simple_mdp.start_all_orders)
+            all_orders=simple_mdp.start_all_orders,
+        )
 
         if counter_drop_forbidden:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP],
-                                         [self.SOUP_DELIVER]
-                                         )
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP],
+                [self.SOUP_DELIVER],
+            )
         else:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP],
-                                         [self.COUNTER_DROP, self.SOUP_DELIVER]
-                                         )
-
-    def simple_mdp_pickup_counter_soup(self, planner, counter_drop_forbidden=False):
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP],
+                [self.COUNTER_DROP, self.SOUP_DELIVER],
+            )
+
+    def simple_mdp_pickup_counter_soup(
+        self, planner, counter_drop_forbidden=False
+    ):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n)],
+            [P((2, 2), n), P((2, 1), n)],
             {(0, 2): done_soup_obj((0, 2))},
-            all_orders=simple_mdp.start_all_orders)
-        self.check_ml_action_manager(s, planner,
-                                     [self.ONION_PICKUP, self.DISH_PICKUP, self.COUNTER_PICKUP],
-                                     [self.ONION_PICKUP, self.DISH_PICKUP, self.COUNTER_PICKUP]
-                                     )
-
-    def simple_mdp_pickup_counter_dish(self, planner, counter_drop_forbidden=False):
+            all_orders=simple_mdp.start_all_orders,
+        )
+        self.check_ml_action_manager(
+            s,
+            planner,
+            [self.ONION_PICKUP, self.DISH_PICKUP, self.COUNTER_PICKUP],
+            [self.ONION_PICKUP, self.DISH_PICKUP, self.COUNTER_PICKUP],
+        )
+
+    def simple_mdp_pickup_counter_dish(
+        self, planner, counter_drop_forbidden=False
+    ):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n)],
-            {(0, 2): Obj('dish', (0, 2))},
-            all_orders=simple_mdp.start_all_orders)
-        self.check_ml_action_manager(s, planner,
-                                     [self.ONION_PICKUP, self.DISH_PICKUP, self.COUNTER_PICKUP],
-                                     [self.ONION_PICKUP, self.DISH_PICKUP, self.COUNTER_PICKUP]
-                                     )
-
-    def simple_mdp_pickup_counter_onion(self, planner, counter_drop_forbidden=False):
+            [P((2, 2), n), P((2, 1), n)],
+            {(0, 2): Obj("dish", (0, 2))},
+            all_orders=simple_mdp.start_all_orders,
+        )
+        self.check_ml_action_manager(
+            s,
+            planner,
+            [self.ONION_PICKUP, self.DISH_PICKUP, self.COUNTER_PICKUP],
+            [self.ONION_PICKUP, self.DISH_PICKUP, self.COUNTER_PICKUP],
+        )
+
+    def simple_mdp_pickup_counter_onion(
+        self, planner, counter_drop_forbidden=False
+    ):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n)],
-            {(0, 2): Obj('onion', (0, 2))},
-            all_orders=simple_mdp.start_all_orders)
-        self.check_ml_action_manager(s, planner,
-                                     [self.ONION_PICKUP, self.DISH_PICKUP, self.COUNTER_PICKUP],
-                                     [self.ONION_PICKUP, self.DISH_PICKUP, self.COUNTER_PICKUP]
-                                     )
-
-    def simple_mdp_drop_useless_dish_with_soup_idle(self, planner, counter_drop_forbidden=False):
+            [P((2, 2), n), P((2, 1), n)],
+            {(0, 2): Obj("onion", (0, 2))},
+            all_orders=simple_mdp.start_all_orders,
+        )
+        self.check_ml_action_manager(
+            s,
+            planner,
+            [self.ONION_PICKUP, self.DISH_PICKUP, self.COUNTER_PICKUP],
+            [self.ONION_PICKUP, self.DISH_PICKUP, self.COUNTER_PICKUP],
+        )
+
+    def simple_mdp_drop_useless_dish_with_soup_idle(
+        self, planner, counter_drop_forbidden=False
+    ):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n, Obj('dish', (2, 1)))],
+            [P((2, 2), n), P((2, 1), n, Obj("dish", (2, 1)))],
             {(2, 0): idle_soup_obj((2, 0), 3)},
-            all_orders=simple_mdp.start_all_orders)
+            all_orders=simple_mdp.start_all_orders,
+        )
         if counter_drop_forbidden:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
-                                         []
-                                         )
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+                [],
+            )
         else:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
-                                         [self.COUNTER_DROP]
-                                         )
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+                [self.COUNTER_DROP],
+            )
 
     def simple_mdp_pickup_soup(self, planner, counter_drop_forbidden=False):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n, Obj('dish', (2, 1)))],
+            [P((2, 2), n), P((2, 1), n, Obj("dish", (2, 1)))],
             {(2, 0): done_soup_obj((2, 0))},
-            all_orders=simple_mdp.start_all_orders)
+            all_orders=simple_mdp.start_all_orders,
+        )
         if counter_drop_forbidden:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP],
-                                         [self.POT_INTERACT]
-                                         )
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP],
+                [self.POT_INTERACT],
+            )
         else:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP],
-                                         [self.COUNTER_DROP, self.POT_INTERACT]
-                                         )
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP],
+                [self.COUNTER_DROP, self.POT_INTERACT],
+            )
 
     def simple_mdp_pickup_dish(self, planner, counter_drop_forbidden=False):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n)],
+            [P((2, 2), n), P((2, 1), n)],
             {(2, 0): done_soup_obj((2, 0))},
-            all_orders=simple_mdp.start_all_orders)
-        self.check_ml_action_manager(s, planner,
-                                     [self.ONION_PICKUP, self.DISH_PICKUP],
-                                     [self.ONION_PICKUP, self.DISH_PICKUP]
-                                     )
-
-    def simple_mdp_start_good_soup_cooking(self, planner, counter_drop_forbidden=False):
+            all_orders=simple_mdp.start_all_orders,
+        )
+        self.check_ml_action_manager(
+            s,
+            planner,
+            [self.ONION_PICKUP, self.DISH_PICKUP],
+            [self.ONION_PICKUP, self.DISH_PICKUP],
+        )
+
+    def simple_mdp_start_good_soup_cooking(
+        self, planner, counter_drop_forbidden=False
+    ):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n)],
+            [P((2, 2), n), P((2, 1), n)],
             {(2, 0): idle_soup_obj((2, 0), 3)},
-            all_orders=simple_mdp.start_all_orders)
-        self.check_ml_action_manager(s, planner,
-                                     [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
-                                     [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT]
-                                     )
-
-    def simple_mdp_start_bad_soup_cooking(self, planner, counter_drop_forbidden=False):
+            all_orders=simple_mdp.start_all_orders,
+        )
+        self.check_ml_action_manager(
+            s,
+            planner,
+            [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+            [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+        )
+
+    def simple_mdp_start_bad_soup_cooking(
+        self, planner, counter_drop_forbidden=False
+    ):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n)],
+            [P((2, 2), n), P((2, 1), n)],
             {(2, 0): idle_soup_obj((2, 0), 2)},
-            all_orders=simple_mdp.start_all_orders)
-        self.check_ml_action_manager(s, planner,
-                                     [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
-                                     [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT]
-                                     )
-
-    def simple_mdp_start_1_onion_soup_cooking(self, planner, counter_drop_forbidden=False):
+            all_orders=simple_mdp.start_all_orders,
+        )
+        self.check_ml_action_manager(
+            s,
+            planner,
+            [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+            [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+        )
+
+    def simple_mdp_start_1_onion_soup_cooking(
+        self, planner, counter_drop_forbidden=False
+    ):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n)],
+            [P((2, 2), n), P((2, 1), n)],
             {(2, 0): idle_soup_obj((2, 0), 1)},
-            all_orders=simple_mdp.start_all_orders)
-        self.check_ml_action_manager(s, planner,
-                                     [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
-                                     [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT]
-                                     )
-
-    def simple_mdp_drop_useless_onion_good_soup(self, planner, counter_drop_forbidden=False):
+            all_orders=simple_mdp.start_all_orders,
+        )
+        self.check_ml_action_manager(
+            s,
+            planner,
+            [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+            [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+        )
+
+    def simple_mdp_drop_useless_onion_good_soup(
+        self, planner, counter_drop_forbidden=False
+    ):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n, Obj('onion', (2, 1)))],
+            [P((2, 2), n), P((2, 1), n, Obj("onion", (2, 1)))],
             {(2, 0): done_soup_obj((2, 0))},
-            all_orders=simple_mdp.start_all_orders)
+            all_orders=simple_mdp.start_all_orders,
+        )
         if counter_drop_forbidden:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP],
-                                         []
-                                         )
+            self.check_ml_action_manager(
+                s, planner, [self.ONION_PICKUP, self.DISH_PICKUP], []
+            )
         else:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP],
-                                         [self.COUNTER_DROP]
-                                         )
-
-    def simple_mdp_drop_useless_onion_bad_soup(self, planner, counter_drop_forbidden=False):
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP],
+                [self.COUNTER_DROP],
+            )
+
+    def simple_mdp_drop_useless_onion_bad_soup(
+        self, planner, counter_drop_forbidden=False
+    ):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n, Obj('onion', (2, 1)))],
+            [P((2, 2), n), P((2, 1), n, Obj("onion", (2, 1)))],
             {(2, 0): done_soup_obj((2, 0), 2)},
-            all_orders=simple_mdp.start_all_orders)
+            all_orders=simple_mdp.start_all_orders,
+        )
         if counter_drop_forbidden:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP],
-                                         []
-                                         )
+            self.check_ml_action_manager(
+                s, planner, [self.ONION_PICKUP, self.DISH_PICKUP], []
+            )
         else:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP],
-                                         [self.COUNTER_DROP]
-                                         )
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP],
+                [self.COUNTER_DROP],
+            )
 
     def simple_mdp_add_3rd_onion(self, planner, counter_drop_forbidden=False):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n, Obj('onion', (2, 1)))],
+            [P((2, 2), n), P((2, 1), n, Obj("onion", (2, 1)))],
             {(2, 0): idle_soup_obj((2, 0), 2)},
-            all_orders=simple_mdp.start_all_orders)
+            all_orders=simple_mdp.start_all_orders,
+        )
         if counter_drop_forbidden:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
-                                         [self.POT_INTERACT]
-                                         )
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+                [self.POT_INTERACT],
+            )
         else:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
-                                         [self.COUNTER_DROP, self.POT_INTERACT]
-                                         )
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+                [self.COUNTER_DROP, self.POT_INTERACT],
+            )
 
     def simple_mdp_add_2nd_onion(self, planner, counter_drop_forbidden=False):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n, Obj('onion', (2, 1)))],
+            [P((2, 2), n), P((2, 1), n, Obj("onion", (2, 1)))],
             {(2, 0): idle_soup_obj((2, 0), 1)},
-            all_orders=simple_mdp.start_all_orders)
+            all_orders=simple_mdp.start_all_orders,
+        )
         if counter_drop_forbidden:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
-                                         [self.POT_INTERACT]
-                                         )
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+                [self.POT_INTERACT],
+            )
         else:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
-                                         [self.COUNTER_DROP, self.POT_INTERACT]
-                                         )
-
-    def simple_mdp_drop_useless_dish(self, planner, counter_drop_forbidden=False):
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+                [self.COUNTER_DROP, self.POT_INTERACT],
+            )
+
+    def simple_mdp_drop_useless_dish(
+        self, planner, counter_drop_forbidden=False
+    ):
         s = OvercookedState(
-            [P((2, 2), n),
-             P((2, 1), n, Obj('dish', (2, 1)))],
+            [P((2, 2), n), P((2, 1), n, Obj("dish", (2, 1)))],
             {(2, 0): idle_soup_obj((2, 0), 1)},
-            all_orders=simple_mdp.start_all_orders)
+            all_orders=simple_mdp.start_all_orders,
+        )
         if counter_drop_forbidden:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
-                                         [self.POT_INTERACT]
-                                         )
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+                [self.POT_INTERACT],
+            )
         else:
-            self.check_ml_action_manager(s, planner,
-                                         [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
-                                         [self.COUNTER_DROP, self.POT_INTERACT]
-                                         )
-
-    def check_ml_action_manager(self, state, am, expected_mla_0, expected_mla_1, debug=False):
+            self.check_ml_action_manager(
+                s,
+                planner,
+                [self.ONION_PICKUP, self.DISH_PICKUP, self.POT_INTERACT],
+                [self.COUNTER_DROP, self.POT_INTERACT],
+            )
+
+    def check_ml_action_manager(
+        self, state, am, expected_mla_0, expected_mla_1, debug=False
+    ):
         """
         args:
             state (OvercookedState): an overcooked state
@@ -630,34 +795,51 @@ def check_ml_action_manager(self, state, am, expected_mla_0, expected_mla_1, deb
         mla_0 = am.get_medium_level_actions(state, player_0)
         mla_1 = am.get_medium_level_actions(state, player_1)
 
-
         if debug:
             print("Player 0 mla", mla_0)
             print("Player 1 mla", mla_1)
             print(am.mdp.state_string(state))
 
-        self.assertEqual(set(mla_0), set(expected_mla_0),
-                         "player 0's ml_action should be " + str(expected_mla_0) +
-                         " but get " + str(mla_0))
-        self.assertEqual(set(mla_1), set(expected_mla_1),
-                         "player 0's ml_action should be " + str(expected_mla_1) +
-                         " but get " + str(mla_1))
+        self.assertEqual(
+            set(mla_0),
+            set(expected_mla_0),
+            "player 0's ml_action should be "
+            + str(expected_mla_0)
+            + " but get "
+            + str(mla_0),
+        )
+        self.assertEqual(
+            set(mla_1),
+            set(expected_mla_1),
+            "player 0's ml_action should be "
+            + str(expected_mla_1)
+            + " but get "
+            + str(mla_1),
+        )
+
 
 class TestScenarios(unittest.TestCase):
     def repetative_runs(self, evaluator, num_games=10):
-        trajectory_0 = evaluator.evaluate_human_model_pair(num_games=num_games, native_eval=True)
-        trajectory_1 = evaluator.evaluate_human_model_pair(num_games=num_games, native_eval=True)
+        trajectory_0 = evaluator.evaluate_human_model_pair(
+            num_games=num_games, native_eval=True
+        )
+        trajectory_1 = evaluator.evaluate_human_model_pair(
+            num_games=num_games, native_eval=True
+        )
 
         h0 = GreedyHumanModel(evaluator.env.mlam)
         h1 = GreedyHumanModel(evaluator.env.mlam)
         ap_hh_2 = AgentPair(h0, h1)
-        trajectory_2 = evaluator.evaluate_agent_pair(agent_pair=ap_hh_2, num_games=num_games, native_eval=True)
+        trajectory_2 = evaluator.evaluate_agent_pair(
+            agent_pair=ap_hh_2, num_games=num_games, native_eval=True
+        )
 
         h3 = GreedyHumanModel(evaluator.env.mlam)
         h4 = GreedyHumanModel(evaluator.env.mlam)
         ap_hh_3 = AgentPair(h3, h4)
-        trajectory_3 = evaluator.evaluate_agent_pair(agent_pair=ap_hh_3, num_games=num_games, native_eval=True)
-
+        trajectory_3 = evaluator.evaluate_agent_pair(
+            agent_pair=ap_hh_3, num_games=num_games, native_eval=True
+        )
 
     def test_scenario_3_no_counter(self):
         # Asymmetric advantage scenario
@@ -676,11 +858,12 @@ def test_scenario_3_no_counter(self):
         start_state = mdp.get_standard_start_state()
 
         env_params = {"start_state_fn": lambda: start_state, "horizon": 1000}
-        eva = AgentEvaluator.from_layout_name(mdp_params, env_params, force_compute=force_compute)
+        eva = AgentEvaluator.from_layout_name(
+            mdp_params, env_params, force_compute=force_compute
+        )
 
         self.repetative_runs(eva)
 
-
     def test_scenario_3_yes_counter(self):
         # Asymmetric advantage scenario
         #
@@ -699,18 +882,25 @@ def test_scenario_3_yes_counter(self):
 
         valid_counters = [(5, 3)]
         one_counter_params = {
-            'start_orientations': False,
-            'wait_allowed': False,
-            'counter_goals': valid_counters,
-            'counter_drop': valid_counters,
-            'counter_pickup': [],
-            'same_motion_goals': True
+            "start_orientations": False,
+            "wait_allowed": False,
+            "counter_goals": valid_counters,
+            "counter_drop": valid_counters,
+            "counter_pickup": [],
+            "same_motion_goals": True,
         }
 
         env_params = {"start_state_fn": lambda: start_state, "horizon": 1000}
-        eva = AgentEvaluator.from_layout_name(mdp_params, env_params, mlam_params=one_counter_params, force_compute=force_compute)
+        eva = AgentEvaluator.from_layout_name(
+            mdp_params,
+            env_params,
+            mlam_params=one_counter_params,
+            force_compute=force_compute,
+        )
 
         self.repetative_runs(eva)
+
+
 # # Deprecated. because of Heuristic
 # class TestHighLevelPlanner(unittest.TestCase):
 #     """The HighLevelPlanner class has been mostly discontinued"""
@@ -736,5 +926,5 @@ def test_scenario_3_yes_counter(self):
 #         # ml_planner_large.get_low_level_action_plan(s, heuristic.simple_heuristic)
 #         # ml_planner_large.get_low_level_action_plan(s, heuristic.hard_heuristic)
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/testing/utils.py b/testing/utils.py
index 648ef15d..6d7c1178 100644
--- a/testing/utils.py
+++ b/testing/utils.py
@@ -1,9 +1,11 @@
 import numpy as np
+
 from overcooked_ai_py.agents.benchmarking import AgentEvaluator
 
 # The point of this function is to generate serialized trajectories for MDP dynamics consistency testing
 # NOTE: If intentionally updating MDP dynamics, this function should be used
 
+
 def generate_serialized_trajectory(mdp, save_path):
     # Saving trajectory for dynamics consistency test
     seed = 0
@@ -15,4 +17,4 @@ def generate_serialized_trajectory(mdp, save_path):
         sparse_reward = np.mean(test_trajs["ep_returns"])
         seed += 1
 
-    AgentEvaluator.save_traj_as_json(test_trajs, save_path)
\ No newline at end of file
+    AgentEvaluator.save_traj_as_json(test_trajs, save_path)
diff --git a/testing/visualization_test.py b/testing/visualization_test.py
index b09f568e..80f16f99 100644
--- a/testing/visualization_test.py
+++ b/testing/visualization_test.py
@@ -1,54 +1,101 @@
-import unittest, os, pygame, copy, json
+import copy
+import json
+import os
+import unittest
+
 import numpy as np
-from overcooked_ai_py.agents.benchmarking import AgentEvaluator
+import pygame
+
 from overcooked_ai_py.agents.agent import RandomAgent
-from overcooked_ai_py.mdp.overcooked_mdp import OvercookedGridworld, OvercookedState
-from overcooked_ai_py.visualization.state_visualizer import StateVisualizer
-from overcooked_ai_py.utils import load_from_json, generate_temporary_file_path
-from overcooked_ai_py.mdp.overcooked_mdp import Recipe
+from overcooked_ai_py.agents.benchmarking import AgentEvaluator
+from overcooked_ai_py.mdp.overcooked_mdp import (
+    OvercookedGridworld,
+    OvercookedState,
+    Recipe,
+)
 from overcooked_ai_py.static import TESTING_DATA_DIR
+from overcooked_ai_py.utils import generate_temporary_file_path, load_from_json
+from overcooked_ai_py.visualization.state_visualizer import StateVisualizer
+
 
 def get_file_count(directory_path):
     path, dirs, files = next(os.walk(directory_path))
     return len(files)
 
+
 state_visualizer_dir = os.path.join(TESTING_DATA_DIR, "test_state_visualizer")
-example_img_path = generate_temporary_file_path(prefix="overcooked_visualized_state_", extension=".png")
+example_img_path = generate_temporary_file_path(
+    prefix="overcooked_visualized_state_", extension=".png"
+)
 
 
 def test_render_state_from_dict(test_dict):
     input_dict = copy.deepcopy(test_dict)
     test_dict = copy.deepcopy(test_dict)
-    test_dict["kwargs"]["state"] = OvercookedState.from_dict(test_dict["kwargs"]["state"])
+    test_dict["kwargs"]["state"] = OvercookedState.from_dict(
+        test_dict["kwargs"]["state"]
+    )
     # check only if it raise error or not, for image fidelity render_state check is used
-    StateVisualizer(**test_dict["config"]).display_rendered_state(img_path=example_img_path, **test_dict["kwargs"])
+    StateVisualizer(**test_dict["config"]).display_rendered_state(
+        img_path=example_img_path, **test_dict["kwargs"]
+    )
 
-    actual_result = pygame.surfarray.array3d(StateVisualizer(**test_dict["config"]).render_state(**test_dict["kwargs"]))
-    expected_result = np.load(os.path.join(state_visualizer_dir, test_dict["result_array_filename"]))
+    actual_result = pygame.surfarray.array3d(
+        StateVisualizer(**test_dict["config"]).render_state(
+            **test_dict["kwargs"]
+        )
+    )
+    expected_result = np.load(
+        os.path.join(state_visualizer_dir, test_dict["result_array_filename"])
+    )
     if not actual_result.shape == expected_result.shape:
         print("test with: ", input_dict["result_array_filename"], "is failed")
-        print("test not passed, wrong output shape", actual_result.shape, "!=", expected_result.shape)
+        print(
+            "test not passed, wrong output shape",
+            actual_result.shape,
+            "!=",
+            expected_result.shape,
+        )
         print(json.dumps(input_dict, indent=4, sort_keys=True))
         return False
 
-    wrong_rows, wrong_columns, wrong_color_channels = np.where(actual_result != expected_result)
-    wrong_coordinates = set([(row, col) for row, col in zip(wrong_rows, wrong_columns)])
+    wrong_rows, wrong_columns, wrong_color_channels = np.where(
+        actual_result != expected_result
+    )
+    wrong_coordinates = set(
+        [(row, col) for row, col in zip(wrong_rows, wrong_columns)]
+    )
     incorrect_pixels_num = len(wrong_coordinates)
-    all_pixels_num = int(expected_result.size/3)
+    all_pixels_num = int(expected_result.size / 3)
     if incorrect_pixels_num:
         wrong_coordinate_list = sorted(list(wrong_coordinates))
         print("test with: ", input_dict["result_array_filename"], "is failed")
-        print("test not passed, wrong color of", incorrect_pixels_num, "pixels out of", all_pixels_num)
-        print("first 100 wrong pixels coordinates", wrong_coordinate_list[:100])
+        print(
+            "test not passed, wrong color of",
+            incorrect_pixels_num,
+            "pixels out of",
+            all_pixels_num,
+        )
+        print(
+            "first 100 wrong pixels coordinates", wrong_coordinate_list[:100]
+        )
         print("coordinate\texpected\tactual")
         for i in range(10):
             (wrong_x, wrong_y) = wrong_coord = wrong_coordinate_list[i]
-            print("%s\t%s\t%s" %(str(wrong_coord), str(expected_result[wrong_x, wrong_y]), str(actual_result[wrong_x, wrong_y])))
+            print(
+                "%s\t%s\t%s"
+                % (
+                    str(wrong_coord),
+                    str(expected_result[wrong_x, wrong_y]),
+                    str(actual_result[wrong_x, wrong_y]),
+                )
+            )
         print("test_dict", json.dumps(input_dict))
         return False
     print("test with: ", input_dict["result_array_filename"], "is ok")
     return True
 
+
 class TestStateVisualizer(unittest.TestCase):
     def setUp(self):
         Recipe.configure({})
@@ -59,7 +106,12 @@ def test_setting_up_configs(self):
         init_config = {"tile_size": 123}
         configure_config = {"tile_size": 234}
         configure_defaults_config = {"tile_size": 345}
-        assert default_values["tile_size"] != init_config["tile_size"] != configure_config["tile_size"] != configure_defaults_config["tile_size"]
+        assert (
+            default_values["tile_size"]
+            != init_config["tile_size"]
+            != configure_config["tile_size"]
+            != configure_defaults_config["tile_size"]
+        )
 
         visualizer = StateVisualizer(**init_config)
         self.assertEqual(init_config["tile_size"], visualizer.tile_size)
@@ -68,69 +120,145 @@ def test_setting_up_configs(self):
         self.assertEqual(configure_config["tile_size"], visualizer.tile_size)
 
         StateVisualizer.configure_defaults(**configure_defaults_config)
-        self.assertEqual(configure_defaults_config["tile_size"], StateVisualizer.DEFAULT_VALUES["tile_size"])
-        self.assertEqual(configure_defaults_config["tile_size"], StateVisualizer().tile_size)
+        self.assertEqual(
+            configure_defaults_config["tile_size"],
+            StateVisualizer.DEFAULT_VALUES["tile_size"],
+        )
+        self.assertEqual(
+            configure_defaults_config["tile_size"], StateVisualizer().tile_size
+        )
 
         invalid_kwargs = {"invalid_argument": 123}
-        self.assertRaises(AssertionError,  StateVisualizer, **invalid_kwargs)
-        self.assertRaises(AssertionError,  StateVisualizer.configure_defaults, **invalid_kwargs)
-        self.assertRaises(AssertionError,  visualizer.configure, **invalid_kwargs)
+        self.assertRaises(AssertionError, StateVisualizer, **invalid_kwargs)
+        self.assertRaises(
+            AssertionError,
+            StateVisualizer.configure_defaults,
+            **invalid_kwargs
+        )
+        self.assertRaises(
+            AssertionError, visualizer.configure, **invalid_kwargs
+        )
 
     def test_properties(self):
-        visualizer = StateVisualizer(tile_size=30, hud_interline_size=7, hud_font_size=26)
-        self.assertEqual(visualizer.scale_by_factor,  2)
-        self.assertEqual(visualizer.hud_line_height, 26+7)
+        visualizer = StateVisualizer(
+            tile_size=30, hud_interline_size=7, hud_font_size=26
+        )
+        self.assertEqual(visualizer.scale_by_factor, 2)
+        self.assertEqual(visualizer.hud_line_height, 26 + 7)
 
     def test_hud_display(self):
-        for d in load_from_json(os.path.join(state_visualizer_dir, "render_state_data_test_hud.json")):
+        for d in load_from_json(
+            os.path.join(
+                state_visualizer_dir, "render_state_data_test_hud.json"
+            )
+        ):
             test_render_state_from_dict(d)
 
     def test_differnet_sizes(self):
-        for d in load_from_json(os.path.join(state_visualizer_dir, "render_state_data_test_sizes.json")):
+        for d in load_from_json(
+            os.path.join(
+                state_visualizer_dir, "render_state_data_test_sizes.json"
+            )
+        ):
             test_render_state_from_dict(d)
 
     def test_cooking_timer_display(self):
-        for d in load_from_json(os.path.join(state_visualizer_dir, "render_state_data_test_cooking_display.json")):
+        for d in load_from_json(
+            os.path.join(
+                state_visualizer_dir,
+                "render_state_data_test_cooking_display.json",
+            )
+        ):
             test_render_state_from_dict(d)
 
     def test_various_states(self):
-        for d in load_from_json(os.path.join(state_visualizer_dir, "render_state_data_test_various.json")):
+        for d in load_from_json(
+            os.path.join(
+                state_visualizer_dir, "render_state_data_test_various.json"
+            )
+        ):
             test_render_state_from_dict(d)
 
     def test_generated_layout_states(self):
-        for d in load_from_json(os.path.join(state_visualizer_dir, "render_state_data_test_generated_layout.json")):
+        for d in load_from_json(
+            os.path.join(
+                state_visualizer_dir,
+                "render_state_data_test_generated_layout.json",
+            )
+        ):
             test_render_state_from_dict(d)
 
     def test_default_hud_data_from_trajectories(self):
-        traj_path = os.path.join(TESTING_DATA_DIR, 'test_state_visualizer', 'test_trajectory.json') # NOTE: for test purposes reward is added here despite there was no soup delivery in trajectory
+        traj_path = os.path.join(
+            TESTING_DATA_DIR, "test_state_visualizer", "test_trajectory.json"
+        )  # NOTE: for test purposes reward is added here despite there was no soup delivery in trajectory
         test_trajectory = AgentEvaluator.load_traj_from_json(traj_path)
-        hud_data_path = os.path.join(TESTING_DATA_DIR, 'test_state_visualizer', 'expected_default_hud_data_from_trajectories.json')
+        hud_data_path = os.path.join(
+            TESTING_DATA_DIR,
+            "test_state_visualizer",
+            "expected_default_hud_data_from_trajectories.json",
+        )
         expected_hud_data = load_from_json(hud_data_path)
-        result_hud_data = StateVisualizer().default_hud_data_from_trajectories(test_trajectory)
-        self.assertEqual(json.dumps(result_hud_data, sort_keys=True),
-            json.dumps(expected_hud_data, sort_keys=True))
+        result_hud_data = StateVisualizer().default_hud_data_from_trajectories(
+            test_trajectory
+        )
+        self.assertEqual(
+            json.dumps(result_hud_data, sort_keys=True),
+            json.dumps(expected_hud_data, sort_keys=True),
+        )
 
     def test_action_probs_display(self):
-        for d in load_from_json(os.path.join(state_visualizer_dir, "render_state_data_test_action_probs_display.json")):
+        for d in load_from_json(
+            os.path.join(
+                state_visualizer_dir,
+                "render_state_data_test_action_probs_display.json",
+            )
+        ):
             test_render_state_from_dict(d)
 
     def test_trajectory_visualization(self):
         # we don't have good way to check slider automatically so its mostly test for basic stuff like numer of outputed images, if using method raises error etc.
-        traj_path = os.path.join(TESTING_DATA_DIR, 'test_state_visualizer', 'test_trajectory.json')
+        traj_path = os.path.join(
+            TESTING_DATA_DIR, "test_state_visualizer", "test_trajectory.json"
+        )
         test_trajectory = AgentEvaluator.load_traj_from_json(traj_path)
         expected_images_num = len(test_trajectory["ep_states"][0])
         assert expected_images_num == 10
-        action_probs = [ [RandomAgent(all_actions=True).action(state)[1]["action_probs"]]*2 for state in test_trajectory["ep_states"][0]]
+        action_probs = [
+            [RandomAgent(all_actions=True).action(state)[1]["action_probs"]]
+            * 2
+            for state in test_trajectory["ep_states"][0]
+        ]
 
-        result_img_directory_path = StateVisualizer().display_rendered_trajectory(test_trajectory, action_probs=action_probs, ipython_display=False)
-        self.assertEqual(get_file_count(result_img_directory_path), expected_images_num)
+        result_img_directory_path = (
+            StateVisualizer().display_rendered_trajectory(
+                test_trajectory,
+                action_probs=action_probs,
+                ipython_display=False,
+            )
+        )
+        self.assertEqual(
+            get_file_count(result_img_directory_path), expected_images_num
+        )
 
-        custom_img_directory_path = generate_temporary_file_path(prefix="overcooked_visualized_trajectory", extension="")
-        self.assertNotEqual(custom_img_directory_path, result_img_directory_path)
-        result_img_directory_path = StateVisualizer().display_rendered_trajectory(test_trajectory, img_directory_path=custom_img_directory_path, ipython_display=False)
+        custom_img_directory_path = generate_temporary_file_path(
+            prefix="overcooked_visualized_trajectory", extension=""
+        )
+        self.assertNotEqual(
+            custom_img_directory_path, result_img_directory_path
+        )
+        result_img_directory_path = (
+            StateVisualizer().display_rendered_trajectory(
+                test_trajectory,
+                img_directory_path=custom_img_directory_path,
+                ipython_display=False,
+            )
+        )
         self.assertEqual(custom_img_directory_path, result_img_directory_path)
-        self.assertEqual(get_file_count(result_img_directory_path), expected_images_num)
+        self.assertEqual(
+            get_file_count(result_img_directory_path), expected_images_num
+        )
 
 
-if __name__ == '__main__':
-    unittest.main()
\ No newline at end of file
+if __name__ == "__main__":
+    unittest.main()