diff --git a/predicators/behavior_utils/behavior_utils.py b/predicators/behavior_utils/behavior_utils.py index f91f102785..2ca7d421dd 100644 --- a/predicators/behavior_utils/behavior_utils.py +++ b/predicators/behavior_utils/behavior_utils.py @@ -40,7 +40,7 @@ 'top_cabinet', 'sofa', 'oatmeal', 'chip', 'vegetable_oil', 'sugar', 'cabinet', 'floor', 'pasta', 'sauce', 'electric_refrigerator', 'olive_oil', 'sugar_jar', 'spaghetti_sauce', 'mayonnaise', 'fridge', 'board_game', - 'video_game', 'facsimile' + 'video_game', 'facsimile', 'tray', 'soap', 'rag', 'newspaper', 'sink', 'scrub_brush', 'oven' } PICK_PLACE_OBJECT_TYPES = { 'mineral_water', 'oatmeal', 'blueberry', 'headset', 'jug', 'flank', @@ -191,6 +191,13 @@ 'facsimile', } +SOAKABLE_OBJECT_TYPES = { + "rag", "scrub_brush", +} + +CLEANABLE_OBJECT_TYPES = { + "oven" +} def get_aabb_volume(lo: Array, hi: Array) -> float: """Simple utility function to compute the volume of an aabb. diff --git a/predicators/behavior_utils/option_model_fns.py b/predicators/behavior_utils/option_model_fns.py index c091bb59d2..f1596fcbc7 100644 --- a/predicators/behavior_utils/option_model_fns.py +++ b/predicators/behavior_utils/option_model_fns.py @@ -440,3 +440,53 @@ def toggleOnObjectOptionModel(_init_state: State, env.step(np.zeros(env.action_space.shape)) return toggleOnObjectOptionModel + +def create_soak_option_model( + plan: List[List[float]], _original_orientation: List[List[float]], + obj_to_soak: "URDFObject") -> Callable[[State, "BehaviorEnv"], None]: + """Instantiates and returns a soak option model given a dummy plan.""" + del plan + + def soakObjectOptionModel(_init_state: State, env: "BehaviorEnv") -> None: + logging.info(f"PRIMITIVE: Attempting to soak {obj_to_soak.name}") + if np.linalg.norm( + np.array(obj_to_soak.get_position()) - + np.array(env.robots[0].get_position())) < 2: + if hasattr( + obj_to_soak, + "states") and object_states.Soaked in obj_to_soak.states: + obj_to_soak.states[object_states.Soaked].set_value(True) + else: + logging.info("PRIMITIVE soak failed, cannot be soaked") + else: + logging.info("PRIMITIVE soak failed, too far") + obj_to_soak.force_wakeup() + # Step the simulator to update visuals. + env.step(np.zeros(env.action_space.shape)) + + return soakObjectOptionModel + +def create_clean_stained_option_model( + plan: List[List[float]], _original_orientation: List[List[float]], + obj_to_clean: "URDFObject") -> Callable[[State, "BehaviorEnv"], None]: + """Instantiates and returns a clean stained option model given a dummy plan.""" + del plan + + def cleanStainedObjectOptionModel(_init_state: State, env: "BehaviorEnv") -> None: + logging.info(f"PRIMITIVE: Attempting to clean stained {obj_to_clean.name}") + if np.linalg.norm( + np.array(obj_to_clean.get_position()) - + np.array(env.robots[0].get_position())) < 2: + if hasattr( + obj_to_clean, + "states") and object_states.Stained in obj_to_clean.states: + obj_to_clean.states[object_states.Stained].set_value(True) + else: + logging.info("PRIMITIVE cleaning failed, cannot be cleaned") + else: + logging.info("PRIMITIVE cleaning failed, too far") + obj_to_clean.force_wakeup() + # Step the simulator to update visuals. + env.step(np.zeros(env.action_space.shape)) + + return cleanStainedObjectOptionModel diff --git a/predicators/envs/behavior.py b/predicators/envs/behavior.py index 3e2492a2fa..cf350e0fa3 100644 --- a/predicators/envs/behavior.py +++ b/predicators/envs/behavior.py @@ -48,7 +48,7 @@ create_close_option_model, create_grasp_option_model, \ create_navigate_option_model, create_open_option_model, \ create_place_inside_option_model, create_place_option_model, \ - create_toggle_on_option_model + create_toggle_on_option_model, create_soak_option_model, create_clean_stained_option_model from predicators.envs import BaseEnv from predicators.settings import CFG from predicators.structs import Action, Array, GroundAtom, Object, \ @@ -150,6 +150,8 @@ def set_options(self) -> None: create_close_option_model, create_place_inside_option_model, create_toggle_on_option_model, + create_soak_option_model, + create_clean_stained_option_model, ] # name, planner_fn, option_policy_fn, option_model_fn, @@ -167,7 +169,10 @@ def set_options(self) -> None: ("PlaceInside", planner_fns[2], option_policy_fns[3], option_model_fns[5], 3, 1, (-1.0, 1.0)), ("ToggleOn", planner_fns[3], option_policy_fns[3], - option_model_fns[6], 3, 1, (-1.0, 1.0))] + option_model_fns[6], 3, 1, (-1.0, 1.0)), + ("Soak", planner_fns[3], option_policy_fns[3], + option_model_fns[7], 3, 1, (-1.0, 1.0)), + ("CleanStained", planner_fns[3], option_policy_fns[3], option_model_fns[8], 3, 1, (-1.0, 1.0))] self._options: Set[ParameterizedOption] = set() for (name, planner_fn, policy_fn, option_model_fn, param_dim, num_args, parameter_limits) in option_elems: @@ -338,6 +343,13 @@ def _get_tasks(self, np.zeros(self.igibson_behavior_env.action_space.shape)) init_state = self.current_ig_state_to_state(use_test_scene=testing) goal = self._get_task_goal() + #### TODO Kathryn + new_goal = set() + for atom in goal: + if "nextto" not in str(atom): + new_goal.add(atom) + goal = new_goal + #### task = Task(init_state, goal) # If the goal already happens to hold in the init state, then # resample. @@ -362,9 +374,13 @@ def _get_task_goal(self) -> Set[GroundAtom]: if head_expr.terms[0] == 'not': # Currently, the only goals that include 'not' are those that # include 'not open' statements, so turn these into 'closed'. - assert head_expr.terms[1] == 'open' - bddl_name = 'closed' - obj_start_idx = 2 + # assert head_expr.terms[1] == 'open' + if head_expr.terms[1] == 'open': + bddl_name = 'closed' + obj_start_idx = 2 + elif head_expr.terms[1] == 'stained': + bddl_name = 'clean' + obj_start_idx = 2 else: bddl_name = head_expr.terms[0] # untyped obj_start_idx = 1 @@ -443,6 +459,12 @@ def predicates(self) -> Set[Predicate]: ("toggled_on", self._toggled_on_classifier, 1), ("toggled-off", self._toggled_off_classifier, 1), ("toggleable", self._toggleable_classifier, 1), + ("soakable", self._soakable_classifier, 1), + ("dry", self._dry_classifier, 1), + ("clean", self._clean_classifier, 1), + ("stainable", self._stainable_classifier, 1), + ("stained", self._stained_classifier, 1), + ("soaked", self._soaked_classifier, 1), ] for name, classifier, arity in custom_predicate_specs: @@ -614,6 +636,7 @@ def object_to_ig_object(self, obj: Object) -> "ArticulatedObject": # Do not add @functools.lru_cache(maxsize=None) here this will # lead to wrong mappings when we load a different scene def _name_to_ig_object(self, name: str) -> "ArticulatedObject": + # import ipdb; ipdb.set_trace() for ig_obj in self._get_task_relevant_objects(): # Name is extended with sub-type in some behavior tasks if self._ig_object_name(ig_obj).startswith(name): @@ -864,6 +887,81 @@ def _toggleable_classifier(self, ig_obj, "states") and object_states.ToggledOn in ig_obj.states return obj_toggleable + def _soakable_classifier(self, + state: State, + objs: Sequence[Object], + skip_allclose_check: bool = False) -> bool: + self.check_state_closeness_and_load(state, skip_allclose_check) + assert len(objs) == 1 + ig_obj = self.object_to_ig_object(objs[0]) + obj_soakable = hasattr( + ig_obj, "states") and object_states.Soaked in ig_obj.states + return obj_soakable + + def _dry_classifier(self, + state: State, + objs: Sequence[Object], + skip_allclose_check: bool = False) -> bool: + self.check_state_closeness_and_load(state, skip_allclose_check) + assert len(objs) == 1 + ig_obj = self.object_to_ig_object(objs[0]) + obj_soakable = self._soakable_classifier(state, objs) + if obj_soakable: + if not ig_obj.states[object_states.Soaked].get_value(): + return True + return False + + def _stainable_classifier(self, + state: State, + objs: Sequence[Object], + skip_allclose_check: bool = False) -> bool: + self.check_state_closeness_and_load(state, skip_allclose_check) + assert len(objs) == 1 + ig_obj = self.object_to_ig_object(objs[0]) + obj_stainable = hasattr( + ig_obj, "states") and object_states.Stained in ig_obj.states + return obj_stainable + + def _clean_classifier(self, + state: State, + objs: Sequence[Object], + skip_allclose_check: bool = False) -> bool: + self.check_state_closeness_and_load(state, skip_allclose_check) + assert len(objs) == 1 + ig_obj = self.object_to_ig_object(objs[0]) + obj_stainable = self._stainable_classifier(state, objs) + if obj_stainable: + if not ig_obj.states[object_states.Stained].get_value(): + return True + return False + + def _stained_classifier(self, + state: State, + objs: Sequence[Object], + skip_allclose_check: bool = False) -> bool: + self.check_state_closeness_and_load(state, skip_allclose_check) + assert len(objs) == 1 + ig_obj = self.object_to_ig_object(objs[0]) + obj_stainable = self._stainable_classifier(state, objs) + if obj_stainable: + if ig_obj.states[object_states.Stained].get_value(): + return True + return False + + def _soaked_classifier(self, + state: State, + objs: Sequence[Object], + skip_allclose_check: bool = False) -> bool: + self.check_state_closeness_and_load(state, skip_allclose_check) + assert len(objs) == 1 + ig_obj = self.object_to_ig_object(objs[0]) + obj_soakable = self._soakable_classifier(state, objs) + if obj_soakable: + if ig_obj.states[object_states.Soaked].get_value(): + return True + return False + + @staticmethod def _ig_object_name(ig_obj: "ArticulatedObject") -> str: if isinstance(ig_obj, (URDFObject, RoomFloor)): diff --git a/predicators/ground_truth_nsrts.py b/predicators/ground_truth_nsrts.py index 98cf31e0c0..5bd22720af 100644 --- a/predicators/ground_truth_nsrts.py +++ b/predicators/ground_truth_nsrts.py @@ -8,7 +8,7 @@ from predicators.behavior_utils.behavior_utils import OPENABLE_OBJECT_TYPES, \ PICK_PLACE_OBJECT_TYPES, PLACE_INTO_SURFACE_OBJECT_TYPES, \ - PLACE_ONTOP_SURFACE_OBJECT_TYPES, TOGGLEABLE_OBJECT_TYPES, \ + PLACE_ONTOP_SURFACE_OBJECT_TYPES, TOGGLEABLE_OBJECT_TYPES, SOAKABLE_OBJECT_TYPES, CLEANABLE_OBJECT_TYPES, \ sample_navigation_params, sample_place_inside_params, \ sample_place_ontop_params from predicators.envs import get_or_create_env @@ -2912,6 +2912,8 @@ def _get_predicate(base_pred_name: str, op_name_count_open = itertools.count() op_name_count_close = itertools.count() op_name_count_place_inside = itertools.count() + op_name_count_soak = itertools.count() + op_name_count_clean_stained = itertools.count() # Dummy sampler definition. Useful for open and close. def dummy_param_sampler(state: State, goal: Set[GroundAtom], @@ -3374,6 +3376,120 @@ def place_inside_obj_pos_sampler( )) nsrts.add(nsrt) + elif base_option_name == "Soak": + assert len(option_arg_type_names) == 1 + soak_obj_type_name = option_arg_type_names[0] + soak_obj_type = type_name_to_type[soak_obj_type_name] + soak_obj = Variable("?obj", soak_obj_type) + + if soak_obj_type.name not in SOAKABLE_OBJECT_TYPES: + continue + + parameters = [soak_obj] + option_vars = [soak_obj] + preconditions = { + _get_lifted_atom("reachable", [soak_obj]), + _get_lifted_atom("dry", [soak_obj]), + _get_lifted_atom("soakable", [soak_obj]), + _get_lifted_atom("holding", [soak_obj]), + } + add_effects = {_get_lifted_atom("soaked", [soak_obj])} + delete_effects = {_get_lifted_atom("dry", [soak_obj])} + nsrt = NSRT( + f"{option.name}-{next(op_name_count_soak)}", parameters, + preconditions, add_effects, delete_effects, set(), option, + option_vars, lambda s, g, r, o: dummy_param_sampler( + s, + g, + r, + [ + env.object_to_ig_object(o_i) + if isinstance(o_i, Object) else o_i for o_i in o + ], + )) + nsrts.add(nsrt) + + elif base_option_name == "CleanStained": + assert len(option_arg_type_names) == 1 + soaked_obj_type_name = option_arg_type_names[0] + soaked_obj_type = type_name_to_type[soaked_obj_type_name] + soaked_obj = Variable("?obj", soaked_obj_type) + + if soaked_obj_type.name not in SOAKABLE_OBJECT_TYPES: + continue + + for clean_obj_types in sorted(env.task_relevant_types): + if clean_obj_types.name not in CLEANABLE_OBJECT_TYPES or \ + clean_obj_types.name == soaked_obj_type.name: + continue + + clean_obj = Variable("?clean", clean_obj_types) + parameters = [clean_obj, soaked_obj] + option_vars = [clean_obj] + + preconditions = { + _get_lifted_atom("holding", [soaked_obj]), + _get_lifted_atom("reachable", [clean_obj]), + _get_lifted_atom("soaked", [soaked_obj]), + _get_lifted_atom("stained", [clean_obj]), + } + + add_effects = {_get_lifted_atom("clean", [clean_obj])} + delete_effects = {_get_lifted_atom("stained", [clean_obj])} + nsrt = NSRT( + f"{option.name}-{next(op_name_count_clean_stained)}", parameters, + preconditions, add_effects, delete_effects, set(), option, + option_vars, lambda s, g, r, o: dummy_param_sampler( + s, + g, + r, + [ + env.object_to_ig_object(o_i) + if isinstance(o_i, Object) else o_i for o_i in o + ], + )) + nsrts.add(nsrt) + + # assert len(option_arg_type_names) == 1 + # clean_obj_type_name = option_arg_type_names[0] + # clean_obj_type = type_name_to_type[clean_obj_type_name] + # clean_obj = Variable("?obj", clean_obj_type) + + # if clean_obj_type.name not in CLEANABLE_OBJECT_TYPES: + # continue + + # for soaked_obj_types in sorted(env.task_relevant_types): + # if soaked_obj_types.name not in SOAKABLE_OBJECT_TYPES or \ + # clean_obj_type.name == soaked_obj_types.name: + # continue + + # soaked_obj = Variable("?soaked", soaked_obj_types) + # parameters = [clean_obj, soaked_obj] + # option_vars = [clean_obj] + + # preconditions = { + # _get_lifted_atom("holding", [soaked_obj]), + # _get_lifted_atom("reachable", [clean_obj]), + # _get_lifted_atom("soaked", [soaked_obj]), + # _get_lifted_atom("stained", [clean_obj]), + # } + + # add_effects = {_get_lifted_atom("clean", [clean_obj])} + # delete_effects = {_get_lifted_atom("stained", [soaked_obj])} + # nsrt = NSRT( + # f"{option.name}-{next(op_name_count_clean_stained)}", parameters, + # preconditions, add_effects, delete_effects, set(), option, + # option_vars, lambda s, g, r, o: dummy_param_sampler( + # s, + # g, + # r, + # [ + # env.object_to_ig_object(o_i) + # if isinstance(o_i, Object) else o_i for o_i in o + # ], + # )) + # nsrts.add(nsrt) + else: raise ValueError( f"Unexpected base option name: {base_option_name}")