diff --git a/src/overcooked_ai_py/mdp/layout_generator.py b/src/overcooked_ai_py/mdp/layout_generator.py index 0b109d44..80474486 100644 --- a/src/overcooked_ai_py/mdp/layout_generator.py +++ b/src/overcooked_ai_py/mdp/layout_generator.py @@ -128,47 +128,49 @@ 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"] + # 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) 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) - if "start_all_orders" in mdp_gen_params: - 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"] - else: - recipe_params = LayoutGenerator.add_generated_mdp_params_orders(self.mdp_params) - if "feature_types" not in mdp_gen_params: mdp_gen_params["feature_types"] = DEFAULT_FEATURE_TYPES - 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, - feature_types=mdp_gen_params["feature_types"], - display=mdp_gen_params["display"] - ) - + 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 + 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"): @@ -199,10 +201,18 @@ def padded_mdp(self, mdp, display=False): start_positions = self.get_random_starting_positions(padded_grid) 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): + 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"] + ) - mdp_params = LayoutGenerator.add_generated_mdp_params_orders(self.mdp_params) - return OvercookedGridworld.from_grid(mdp_grid, base_layout_params=mdp_params) - 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) diff --git a/testing/overcooked_test.py b/testing/overcooked_test.py index 9713016a..542bc4d4 100644 --- a/testing/overcooked_test.py +++ b/testing/overcooked_test.py @@ -899,9 +899,15 @@ def test_random_layout_feature_types(self): 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": (1, 1), - "feature_types": used_features} - mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(**mdp_gen_params) + 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() @@ -916,31 +922,43 @@ def test_random_layout_generated_recipes(self): only_onions_dict_recipes = [r.to_dict() for r in only_onions_recipes] # checking if recipes are generated from mdp_params - mdp_params = {"generate_all_orders": {"n":2, "ingredients": ["onion"], "min_size":2, "max_size":3}} - mdp_gen_params = {"mdp_params": mdp_params} - mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(**mdp_gen_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.assertTrue(len(env.mdp.start_bonus_orders) == 0) + self.assertEqual(len(env.mdp.start_bonus_orders), 0) # checking if bonus_orders is subset of all_orders even if not specified - mdp_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}} - mdp_gen_params = {"mdp_params": mdp_params} - mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(**mdp_gen_params) + + 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.assertTrue(len(env.mdp.start_bonus_orders) == 1) + 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_params = {"generate_all_orders": {"n":3, "min_size":2, "max_size":3}} - mdp_gen_params = {"mdp_params": mdp_params} - mdp_fn = LayoutGenerator.mdp_gen_fn_from_dict(**mdp_gen_params) + 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):