Merge pull request #41 from bmielnicki/recipes_generation

Recipes generation

Author: mesutyang97

src/overcooked_ai_py/mdp/layout_generator.py

@@ -1,8 +1,9 @@
 import numpy as np
-import random
+
+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
+from overcooked_ai_py.mdp.overcooked_mdp import OvercookedGridworld, Recipe
 
 EMPTY = ' '
 COUNTER = 'X'
@@ -81,6 +82,7 @@ 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
 
 class LayoutGenerator(object):
     # NOTE: This class hasn't been tested extensively.
@@ -126,36 +128,71 @@ 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)
+
         outer_shape = self.outer_shape
         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", "start_all_orders"]
+            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"): 
+                required_keys.append("start_all_orders")
             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)
-
-            recipe_params = {"start_all_orders": mdp_gen_params["start_all_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"]
-
             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)
-            mdp_generator_fn = lambda: layout_generator.make_disjoint_sets_layout(
-                inner_shape=mdp_gen_params["inner_shape"],
-                prop_empty=mdp_gen_params["prop_empty"],
-                prop_features=mdp_gen_params["prop_feats"],
-                base_param=recipe_params,
-                display=mdp_gen_params["display"]
-            )
-
+            
+            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)
         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"]}
+        if mdp_gen_params.get("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):
+        """
+        adds generated parameters (i.e. generated orders) to mdp_params,
+        returns onchanged copy of mdp_params when there is no "generate_all_orders" and "generate_bonus_orders" keys inside 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"])
+
+            if all_orders_kwargs.get("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:
+            all_recipes = Recipe.ALL_RECIPES
+
+        if mdp_params.get("generate_bonus_orders"):
+            bonus_orders_kwargs = copy.deepcopy(mdp_params["generate_bonus_orders"])
+
+            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]
+        return mdp_params
 
     def padded_mdp(self, mdp, display=False):
         """Returns a padded MDP from an MDP"""
@@ -166,10 +203,20 @@ def padded_mdp(self, mdp, display=False):
         mdp_grid = self.padded_grid_to_layout_grid(padded_grid, start_positions, display=display)
         return OvercookedGridworld.from_grid(mdp_grid)
 
-    def make_disjoint_sets_layout(self, inner_shape, prop_empty, prop_features, base_param={}, display=True):
+    def make_new_layout(self, mdp_gen_params):
+        return self.make_disjoint_sets_layout(
+            inner_shape=mdp_gen_params["inner_shape"],
+            prop_empty=mdp_gen_params["prop_empty"],
+            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):        
         grid = Grid(inner_shape)
         self.dig_space_with_disjoint_sets(grid, prop_empty)
-        self.add_features(grid, prop_features)
+        self.add_features(grid, prop_features, feature_types)
 
         padded_grid = self.embed_grid(grid)
         start_positions = self.get_random_starting_positions(padded_grid)
@@ -242,12 +289,10 @@ 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):
+    def add_features(self, grid, prop_features=0, feature_types=DEFAULT_FEATURE_TYPES):
         """
         Places one round of basic features and then adds random features 
         until prop_features of valid locations are filled"""
-        feature_types = [POT, ONION_DISPENSER, DISH_DISPENSER, SERVING_LOC]
-        # feature_types = [POT, ONION_DISPENSER, TOMATO_DISPENSER, DISH_DISPENSER, SERVING_LOC]  # NOTE: currently disabled TOMATO_DISPENSER
 
         valid_locations = grid.valid_feature_locations()
         np.random.shuffle(valid_locations)

src/overcooked_ai_py/mdp/overcooked_mdp.py

@@ -236,6 +236,34 @@ def configure(cls, conf):
 
         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):
+        """
+        n (int): how many recipes generate
+        min_size (int): min generated recipe size
+        max_size (int): max generated recipe size
+        ingredients (list(str)): list of ingredients used for generating recipes (default is cls.ALL_INGREDIENTS)
+        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
+
+        ingredients = set(ingredients or cls.ALL_INGREDIENTS)
+        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)
+
+        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)]
+        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):

testing/overcooked_test.py

@@ -1,17 +1,17 @@
 import unittest, os
+import json
 import numpy as np
 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.layout_generator import LayoutGenerator
+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
 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 utils import TESTING_DATA_DIR, generate_serialized_trajectory
 
-
 START_ORDER_LIST = ["any"]
 n, s = Direction.NORTH, Direction.SOUTH
 e, w = Direction.EAST, Direction.WEST
@@ -80,6 +80,24 @@ def test_invalid_input(self):
         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.assertEqual(len(Recipe.generate_random_recipes(n=3)), 3)
+        self.assertEqual(len(Recipe.generate_random_recipes(n=99, unique=False)), 99)
+
+        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"])]
+        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 
 
     def _expected_num_recipes(self, num_ingredients, max_len):
         return comb(num_ingredients + max_len, num_ingredients) - 1
@@ -873,6 +891,81 @@ def test_random_layout(self):
         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}]
+
+        for optional_features_combo in optional_features_combinations:
+            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))
+            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
+                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
+
+    def test_random_layout_generated_recipes(self):
+        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))
+        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.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))
+        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.assertEqual(len(env.mdp.start_bonus_orders), 1)
+            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))
+        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}
+        self.assertTrue(len(generated_recipes_strings) > 3)
+        
 
 class TestGymEnvironment(unittest.TestCase):