yapf

Author: yichao-liang

predicators/approaches/__init__.py

@@ -8,8 +8,8 @@
 from predicators import utils
 from predicators.approaches.base_approach import ApproachFailure, \
     ApproachTimeout, BaseApproach, BaseApproachWrapper
-from predicators.structs import ParameterizedOption, Predicate, Task, Type,\
-    ConceptPredicate
+from predicators.structs import ConceptPredicate, ParameterizedOption, \
+    Predicate, Task, Type
 
 __all__ = ["BaseApproach", "ApproachTimeout", "ApproachFailure"]
 
@@ -47,14 +47,21 @@ def create_approach(name: str, initial_predicates: Set[Predicate],
                                         train_tasks)
         # Find wrapper.
         wrapper_cls = _get_wrapper_cls_from_name(wrapper_name)
-        return wrapper_cls(base_approach, initial_predicates, 
-                           initial_options,
-                           types, action_space, train_tasks,
-                           )
+        return wrapper_cls(
+            base_approach,
+            initial_predicates,
+            initial_options,
+            types,
+            action_space,
+            train_tasks,
+        )
 
     # Handle main approaches.
     cls = _get_approach_cls_from_name(name)
-    return cls(initial_predicates, 
-               initial_options, types, action_space,
-               train_tasks,
-               )
+    return cls(
+        initial_predicates,
+        initial_options,
+        types,
+        action_space,
+        train_tasks,
+    )

predicators/approaches/base_approach.py

@@ -8,20 +8,24 @@
 from gym.spaces import Box
 
 from predicators.settings import CFG
-from predicators.structs import Action, Dataset, InteractionRequest, \
-    InteractionResult, Metrics, ParameterizedOption, Predicate, State, Task, \
-    Type, ConceptPredicate
+from predicators.structs import Action, ConceptPredicate, Dataset, \
+    InteractionRequest, InteractionResult, Metrics, ParameterizedOption, \
+    Predicate, State, Task, Type
 from predicators.utils import ExceptionWithInfo
 
 
 class BaseApproach(abc.ABC):
     """Base approach."""
 
-    def __init__(self, initial_predicates: Set[Predicate],
-                 initial_options: Set[ParameterizedOption], types: Set[Type],
-                 action_space: Box, train_tasks: List[Task],
-                 initial_concept_predicates: Set[ConceptPredicate],
-                 ) -> None:
+    def __init__(
+        self,
+        initial_predicates: Set[Predicate],
+        initial_options: Set[ParameterizedOption],
+        types: Set[Type],
+        action_space: Box,
+        train_tasks: List[Task],
+        initial_concept_predicates: Set[ConceptPredicate],
+    ) -> None:
         """All approaches are initialized with only the necessary information
         about the environment."""
         self._initial_predicates = initial_predicates

predicators/approaches/bilevel_planning_approach.py

@@ -17,30 +17,33 @@
 from predicators.planning import PlanningFailure, PlanningTimeout, \
     run_task_plan_once, sesame_plan
 from predicators.settings import CFG
-from predicators.structs import NSRT, Action, GroundAtom, Metrics, \
-    ParameterizedOption, Predicate, State, Task, Type, _GroundNSRT, _Option,\
-    ConceptPredicate
+from predicators.structs import NSRT, Action, ConceptPredicate, GroundAtom, \
+    Metrics, ParameterizedOption, Predicate, State, Task, Type, _GroundNSRT, \
+    _Option
 
 
 class BilevelPlanningApproach(BaseApproach):
     """Bilevel planning approach."""
 
-    def __init__(self,
-                 initial_predicates: Set[Predicate],
-                 initial_options: Set[ParameterizedOption],
-                 types: Set[Type],
-                 action_space: Box,
-                 train_tasks: List[Task],
-                 task_planning_heuristic: str = "default",
-                 max_skeletons_optimized: int = -1,
-                 bilevel_plan_without_sim: Optional[bool] = None,
-                 option_model: Optional[_OptionModelBase] = None,
-                 initial_concept_predicates: Set[ConceptPredicate] = set(),
-                 ) -> None:
-        super().__init__(initial_predicates, 
-                        initial_options, types,
-                        action_space, train_tasks,
-                        initial_concept_predicates=initial_concept_predicates)
+    def __init__(
+            self,
+            initial_predicates: Set[Predicate],
+            initial_options: Set[ParameterizedOption],
+            types: Set[Type],
+            action_space: Box,
+            train_tasks: List[Task],
+            task_planning_heuristic: str = "default",
+            max_skeletons_optimized: int = -1,
+            bilevel_plan_without_sim: Optional[bool] = None,
+            option_model: Optional[_OptionModelBase] = None,
+            initial_concept_predicates: Set[ConceptPredicate] = set(),
+    ) -> None:
+        super().__init__(initial_predicates,
+                         initial_options,
+                         types,
+                         action_space,
+                         train_tasks,
+                         initial_concept_predicates=initial_concept_predicates)
         if task_planning_heuristic == "default":
             task_planning_heuristic = CFG.sesame_task_planning_heuristic
         if max_skeletons_optimized == -1:
@@ -191,12 +194,13 @@ def _get_current_predicates(self) -> Set[Predicate]:
         Defaults to initial predicates.
         """
         return self._initial_predicates | self._initial_concept_predicates
-    
+
     def _get_current_concept_predicates(self) -> Set[ConceptPredicate]:
-        """Get the current set of concept predicates.
-        """
-        cnpt_preds = set([pred for pred in self._get_current_predicates() if
-                          isinstance(pred, ConceptPredicate)])
+        """Get the current set of concept predicates."""
+        cnpt_preds = set([
+            pred for pred in self._get_current_predicates()
+            if isinstance(pred, ConceptPredicate)
+        ])
         return cnpt_preds
 
     def get_option_model(self) -> _OptionModelBase:

predicators/approaches/nsrt_learning_approach.py

@@ -6,8 +6,8 @@
 
 import logging
 import time
-from typing import Any, Dict, List, Optional, Set
 from itertools import chain
+from typing import Any, Dict, List, Optional, Set
 
 import dill as pkl
 from gym.spaces import Box
@@ -18,9 +18,9 @@
 from predicators.nsrt_learning.nsrt_learning_main import learn_nsrts_from_data
 from predicators.planning import task_plan, task_plan_grounding
 from predicators.settings import CFG
-from predicators.structs import NSRT, Dataset, GroundAtomTrajectory, \
-    LowLevelTrajectory, ParameterizedOption, Predicate, Segment, Task, Type,\
-    ConceptPredicate
+from predicators.structs import NSRT, ConceptPredicate, Dataset, \
+    GroundAtomTrajectory, LowLevelTrajectory, ParameterizedOption, Predicate, \
+    Segment, Task, Type
 
 
 class NSRTLearningApproach(BilevelPlanningApproach):
@@ -106,7 +106,7 @@ def _learn_nsrts(self, trajectories: List[LowLevelTrajectory],
         # TODO: make sure it expands until no more new predicates are added
         aux_preds = set(chain.from_iterable(p.auxiliary_concepts for p
                                 in (self._get_current_predicates() |\
-                                self._get_current_concept_predicates()) if 
+                                self._get_current_concept_predicates()) if
                                     isinstance(p, ConceptPredicate) and\
                                     p.auxiliary_concepts))
         self._nsrts, self._segmented_trajs, self._seg_to_nsrt = \

predicators/approaches/oracle_approach.py

@@ -15,26 +15,27 @@
 from predicators.ground_truth_models import get_gt_nsrts
 from predicators.option_model import _OptionModelBase
 from predicators.settings import CFG
-from predicators.structs import NSRT, ParameterizedOption, Predicate, Task, \
-    Type, ConceptPredicate
+from predicators.structs import NSRT, ConceptPredicate, ParameterizedOption, \
+    Predicate, Task, Type
 
 
 class OracleApproach(BilevelPlanningApproach):
     """A bilevel planning approach that uses hand-specified NSRTs."""
 
-    def __init__(self,
-                 initial_predicates: Set[Predicate],
-                 initial_options: Set[ParameterizedOption],
-                 types: Set[Type],
-                 action_space: Box,
-                 train_tasks: List[Task],
-                 task_planning_heuristic: str = "default",
-                 max_skeletons_optimized: int = -1,
-                 bilevel_plan_without_sim: Optional[bool] = None,
-                 nsrts: Optional[Set[NSRT]] = None,
-                 option_model: Optional[_OptionModelBase] = None,
-                 initial_concept_predicates: Set[ConceptPredicate] = set(),
-                 ) -> None:
+    def __init__(
+            self,
+            initial_predicates: Set[Predicate],
+            initial_options: Set[ParameterizedOption],
+            types: Set[Type],
+            action_space: Box,
+            train_tasks: List[Task],
+            task_planning_heuristic: str = "default",
+            max_skeletons_optimized: int = -1,
+            bilevel_plan_without_sim: Optional[bool] = None,
+            nsrts: Optional[Set[NSRT]] = None,
+            option_model: Optional[_OptionModelBase] = None,
+            initial_concept_predicates: Set[ConceptPredicate] = set(),
+    ) -> None:
         super().__init__(initial_predicates,
                          initial_options,
                          types,