diff --git a/predicators/approaches/__init__.py b/predicators/approaches/__init__.py index 502e46736b..ee5ad6fa52 100644 --- a/predicators/approaches/__init__.py +++ b/predicators/approaches/__init__.py @@ -17,17 +17,17 @@ def _get_approach_cls_from_name(name: str) -> TypingType[BaseApproach]: - for cls in utils.get_all_subclasses(BaseApproach): - if not cls.__abstractmethods__ and cls.get_name() == name: - return cls - raise NotImplementedError(f"Unknown approach: {name}") + cls = utils.get_registered_subclass(BaseApproach, name) + if cls is None: + raise NotImplementedError(f"Unknown approach: {name}") + return cls def _get_wrapper_cls_from_name(name: str) -> TypingType[BaseApproachWrapper]: - for cls in utils.get_all_subclasses(BaseApproachWrapper): - if not cls.__abstractmethods__ and cls.get_name() == name: - return cls - raise NotImplementedError(f"Unknown wrapper approach: {name}") + cls = utils.get_registered_subclass(BaseApproachWrapper, name) + if cls is None: + raise NotImplementedError(f"Unknown wrapper approach: {name}") + return cls def create_approach(name: str, initial_predicates: Set[Predicate], diff --git a/predicators/approaches/human_low_level_control_approach.py b/predicators/approaches/human_low_level_control_approach.py index 54ed56a8f2..3d9fb52280 100644 --- a/predicators/approaches/human_low_level_control_approach.py +++ b/predicators/approaches/human_low_level_control_approach.py @@ -354,21 +354,12 @@ def _get_shadow_robot_for_env() -> SingleArmPyBulletRobot: ground plane, and ask the env class for a robot. """ # pylint: disable=import-outside-toplevel - from predicators.envs.base_env import BaseEnv from predicators.envs.pybullet_env import PyBulletEnv # pylint: enable=import-outside-toplevel env_name = CFG.env - env_cls = None - for cls in utils.get_all_subclasses(BaseEnv): - if cls.__abstractmethods__: - continue - if not issubclass(cls, PyBulletEnv): - continue - if cls.get_name() == env_name: - env_cls = cls - break + env_cls = utils.get_registered_subclass(PyBulletEnv, env_name) if env_cls is None: raise NotImplementedError( f"human_low_level_control: no PyBulletEnv subclass registered " diff --git a/predicators/bridge_policies/__init__.py b/predicators/bridge_policies/__init__.py index d73ab03441..de3a3fdf9b 100644 --- a/predicators/bridge_policies/__init__.py +++ b/predicators/bridge_policies/__init__.py @@ -18,10 +18,7 @@ def create_bridge_policy(name: str, types: Set[Type], options: Set[ParameterizedOption], nsrts: Set[NSRT]) -> BaseBridgePolicy: """Create a bridge policy given its name.""" - for cls in utils.get_all_subclasses(BaseBridgePolicy): - if not cls.__abstractmethods__ and cls.get_name() == name: - bridge_policy = cls(types, predicates, options, nsrts) - break - else: + cls = utils.get_registered_subclass(BaseBridgePolicy, name) + if cls is None: raise NotImplementedError(f"Unknown bridge policy: {name}") - return bridge_policy + return cls(types, predicates, options, nsrts) diff --git a/predicators/competence_models.py b/predicators/competence_models.py index 6bd99959fa..1a27f6d2dd 100644 --- a/predicators/competence_models.py +++ b/predicators/competence_models.py @@ -240,10 +240,10 @@ def _run_map_inference(self, betas: List[BetaRV]) -> List[float]: def _get_competence_model_cls_from_name( name: str) -> TypingType[SkillCompetenceModel]: - for cls in utils.get_all_subclasses(SkillCompetenceModel): - if not cls.__abstractmethods__ and cls.get_name() == name: - return cls - raise NotImplementedError(f"Unknown competence model: {name}") + cls = utils.get_registered_subclass(SkillCompetenceModel, name) + if cls is None: + raise NotImplementedError(f"Unknown competence model: {name}") + return cls def create_competence_model(model_name: str, diff --git a/predicators/envs/__init__.py b/predicators/envs/__init__.py index 2510edd608..22a414ca58 100644 --- a/predicators/envs/__init__.py +++ b/predicators/envs/__init__.py @@ -22,12 +22,10 @@ def create_new_env(name: str, If do_cache is True, then cache this env instance so that it can later be loaded using get_or_create_env(). """ - for cls in utils.get_all_subclasses(BaseEnv): - if not cls.__abstractmethods__ and cls.get_name() == name: - env = cls(use_gui, **kwargs) - break - else: + cls = utils.get_registered_subclass(BaseEnv, name) + if cls is None: raise NotImplementedError(f"Unknown env: {name}") + env = cls(use_gui, **kwargs) if do_cache: _MOST_RECENT_ENV_INSTANCE[name] = env return env diff --git a/predicators/envs/pybullet_domino_real.py b/predicators/envs/pybullet_domino_real.py index cb4237e7eb..fb50213c39 100644 --- a/predicators/envs/pybullet_domino_real.py +++ b/predicators/envs/pybullet_domino_real.py @@ -123,11 +123,9 @@ class RealSceneGeometryMixin: perceived scene, while :class:`PyBulletDominoRealGeometryEnv` pairs it with the ordinary generated tasks. - Deliberately NOT a ``BaseEnv`` subclass. ``create_new_env`` resolves an - env by scanning ``get_all_subclasses(BaseEnv)`` for a matching - ``get_name()``, so an intermediate env class here would inherit - ``PyBulletDominoEnv.get_name()`` and shadow the real - ``pybullet_domino``. A plain mixin never enters that scan. + Deliberately NOT a ``BaseEnv`` subclass: it adds geometry to whichever + domino env it is mixed into and has no env name or tasks of its own, + so it stays out of the env registry that ``create_new_env`` reads. """ # Annotations only, no values: these ClassVars belong to the domino env diff --git a/predicators/execution_monitoring/__init__.py b/predicators/execution_monitoring/__init__.py index d73bf94ca3..004ce78097 100644 --- a/predicators/execution_monitoring/__init__.py +++ b/predicators/execution_monitoring/__init__.py @@ -12,10 +12,7 @@ def create_execution_monitor(name: str, ) -> BaseExecutionMonitor: """Create an execution monitor given its name.""" - for cls in utils.get_all_subclasses(BaseExecutionMonitor): - if not cls.__abstractmethods__ and cls.get_name() == name: - execution_monitor = cls() - break - else: + cls = utils.get_registered_subclass(BaseExecutionMonitor, name) + if cls is None: raise NotImplementedError(f"Unrecognized execution monitor: {name}") - return execution_monitor + return cls() diff --git a/predicators/explorers/__init__.py b/predicators/explorers/__init__.py index 3a930dbbe7..924d53a638 100644 --- a/predicators/explorers/__init__.py +++ b/predicators/explorers/__init__.py @@ -47,67 +47,59 @@ def create_explorer( """Create an explorer given its name.""" if max_steps_before_termination is None: max_steps_before_termination = CFG.max_num_steps_interaction_request - for cls in utils.get_all_subclasses(BaseExplorer): - if not cls.__abstractmethods__ and cls.get_name() == name: - # Special case GLIB because it uses babble predicates and an atom - # score function. - if name == "glib": - assert nsrts is not None - assert option_model is not None - assert babble_predicates is not None - assert atom_score_fn is not None - explorer = cls(initial_predicates, initial_options, types, - action_space, train_tasks, - max_steps_before_termination, nsrts, - option_model, babble_predicates, atom_score_fn) - # Special case greedy lookahead because it uses a state score - # function. - elif name == "greedy_lookahead": - assert nsrts is not None - assert option_model is not None - assert state_score_fn is not None - explorer = cls(initial_predicates, initial_options, types, - action_space, train_tasks, - max_steps_before_termination, nsrts, - option_model, state_score_fn) - # Bilevel planning approaches use NSRTs and an option model. - elif issubclass(cls, BilevelPlanningExplorer): - assert nsrts is not None - assert option_model is not None - explorer = cls(initial_predicates, initial_options, types, - action_space, train_tasks, - max_steps_before_termination, nsrts, - option_model) - # Random NSRTs explorer uses NSRTs, but not an option model - elif name == "random_nsrts": - assert nsrts is not None - explorer = cls(initial_predicates, initial_options, types, - action_space, train_tasks, - max_steps_before_termination, nsrts) - # Active sampler explorer uses ground_op_hist and no option model. - elif name == "active_sampler": - assert ground_op_hist is not None - assert competence_models is not None - assert nsrt_to_explorer_sampler is not None - assert seen_train_task_idxs is not None - assert pursue_task_goal_first is not None - explorer = cls(initial_predicates, initial_options, types, - action_space, train_tasks, - max_steps_before_termination, nsrts, - ground_op_hist, competence_models, - nsrt_to_explorer_sampler, seen_train_task_idxs, - pursue_task_goal_first) - elif name == "maple_q": - assert maple_q_function is not None - explorer = cls(initial_predicates, initial_options, types, - action_space, train_tasks, - max_steps_before_termination, nsrts, - maple_q_function) - else: - explorer = cls(initial_predicates, initial_options, types, - action_space, train_tasks, - max_steps_before_termination) - break - else: + cls = utils.get_registered_subclass(BaseExplorer, name) + if cls is None: raise NotImplementedError(f"Unrecognized explorer: {name}") + # Special case GLIB because it uses babble predicates and an atom + # score function. + if name == "glib": + assert nsrts is not None + assert option_model is not None + assert babble_predicates is not None + assert atom_score_fn is not None + explorer = cls(initial_predicates, initial_options, types, + action_space, train_tasks, max_steps_before_termination, + nsrts, option_model, babble_predicates, atom_score_fn) + # Special case greedy lookahead because it uses a state score + # function. + elif name == "greedy_lookahead": + assert nsrts is not None + assert option_model is not None + assert state_score_fn is not None + explorer = cls(initial_predicates, initial_options, types, + action_space, train_tasks, max_steps_before_termination, + nsrts, option_model, state_score_fn) + # Bilevel planning approaches use NSRTs and an option model. + elif issubclass(cls, BilevelPlanningExplorer): + assert nsrts is not None + assert option_model is not None + explorer = cls(initial_predicates, initial_options, types, + action_space, train_tasks, max_steps_before_termination, + nsrts, option_model) + # Random NSRTs explorer uses NSRTs, but not an option model + elif name == "random_nsrts": + assert nsrts is not None + explorer = cls(initial_predicates, initial_options, types, + action_space, train_tasks, max_steps_before_termination, + nsrts) + # Active sampler explorer uses ground_op_hist and no option model. + elif name == "active_sampler": + assert ground_op_hist is not None + assert competence_models is not None + assert nsrt_to_explorer_sampler is not None + assert seen_train_task_idxs is not None + assert pursue_task_goal_first is not None + explorer = cls(initial_predicates, initial_options, types, + action_space, train_tasks, max_steps_before_termination, + nsrts, ground_op_hist, competence_models, + nsrt_to_explorer_sampler, seen_train_task_idxs, + pursue_task_goal_first) + elif name == "maple_q": + assert maple_q_function is not None + explorer = cls(initial_predicates, initial_options, types, + action_space, train_tasks, max_steps_before_termination, + nsrts, maple_q_function) + else: + explorer = cls(initial_predicates, initial_options, types, + action_space, train_tasks, max_steps_before_termination) return explorer diff --git a/predicators/ground_truth_models/domino/options.py b/predicators/ground_truth_models/domino/options.py index fa2308ea74..b0f65fa6c7 100644 --- a/predicators/ground_truth_models/domino/options.py +++ b/predicators/ground_truth_models/domino/options.py @@ -29,11 +29,9 @@ def _skill_robot_env_cls(env_name: str) -> TypingType[PyBulletEnv]: """ # pylint: disable=import-outside-toplevel # local: avoid import cycle from predicators.envs.base_env import BaseEnv - from predicators.utils import get_all_subclasses - for c in get_all_subclasses(BaseEnv): - if not c.__abstractmethods__ and c.get_name() == env_name: - return c # type: ignore[return-value] - return PyBulletDominoEnv + from predicators.utils import get_registered_subclass + cls = get_registered_subclass(BaseEnv, env_name) + return PyBulletDominoEnv if cls is None else cls class PyBulletDominoGroundTruthOptionFactory(_DominoLegacyOptionsMixin, diff --git a/predicators/nsrt_learning/strips_learning/__init__.py b/predicators/nsrt_learning/strips_learning/__init__.py index 84015e2f7e..84bfa771b2 100644 --- a/predicators/nsrt_learning/strips_learning/__init__.py +++ b/predicators/nsrt_learning/strips_learning/__init__.py @@ -28,13 +28,9 @@ def learn_strips_operators(trajectories: List[LowLevelTrajectory], Return a list of PNADs with op (STRIPSOperator), datastore, and option_spec fields filled in (but not sampler). """ - for cls in utils.get_all_subclasses(BaseSTRIPSLearner): - if not cls.__abstractmethods__ and \ - cls.get_name() == CFG.strips_learner: - learner = cls(trajectories, train_tasks, predicates, - segmented_trajs, verify_harmlessness, annotations, - verbose, **kwargs) - break - else: + cls = utils.get_registered_subclass(BaseSTRIPSLearner, CFG.strips_learner) + if cls is None: raise ValueError(f"Unrecognized STRIPS learner: {CFG.strips_learner}") + learner = cls(trajectories, train_tasks, predicates, segmented_trajs, + verify_harmlessness, annotations, verbose, **kwargs) return learner.learn() diff --git a/predicators/perception/__init__.py b/predicators/perception/__init__.py index 7984f70738..edbf3d7af9 100644 --- a/predicators/perception/__init__.py +++ b/predicators/perception/__init__.py @@ -11,10 +11,7 @@ def create_perceiver(name: str, ) -> BasePerceiver: """Create a perceiver given its name.""" - for cls in utils.get_all_subclasses(BasePerceiver): - if not cls.__abstractmethods__ and cls.get_name() == name: - perceiver = cls() - break - else: + cls = utils.get_registered_subclass(BasePerceiver, name) + if cls is None: raise NotImplementedError(f"Unrecognized perceiver: {name}") - return perceiver + return cls() diff --git a/predicators/refinement_estimators/__init__.py b/predicators/refinement_estimators/__init__.py index 501a1d7637..b9f9c9fee9 100644 --- a/predicators/refinement_estimators/__init__.py +++ b/predicators/refinement_estimators/__init__.py @@ -12,10 +12,7 @@ def create_refinement_estimator(name: str) -> BaseRefinementEstimator: """Create an approach given its name.""" - for cls in utils.get_all_subclasses(BaseRefinementEstimator): - if not cls.__abstractmethods__ and cls.get_name() == name: - estimator = cls() - break - else: + cls = utils.get_registered_subclass(BaseRefinementEstimator, name) + if cls is None: raise NotImplementedError(f"Unknown refinement cost estimator: {name}") - return estimator + return cls() diff --git a/predicators/utils.py b/predicators/utils.py index c5a871c504..0de16194ab 100644 --- a/predicators/utils.py +++ b/predicators/utils.py @@ -5162,6 +5162,34 @@ def get_all_subclasses(cls: Any) -> Set[Any]: [s for c in cls.__subclasses__() for s in get_all_subclasses(c)]) +def get_registered_subclass(base: Any, name: str) -> Optional[Any]: + """The concrete subclass of ``base`` registered under ``name``. + + The name belongs to the most general concrete subclass whose + ``get_name()`` returns it. A subclass that inherits the name, such + as a test's instrumented env, never resolves in its place: picking + the first match of ``get_all_subclasses`` would follow set order, + which differs from process to process. Returns None when no concrete + subclass has the name, and raises ValueError when unrelated + subclasses claim it. + """ + matches = [ + cls for cls in get_all_subclasses(base) + if not cls.__abstractmethods__ and cls.get_name() == name + ] + roots = [ + cls for cls in matches + if not any(other is not cls and issubclass(cls, other) + for other in matches) + ] + if len(roots) > 1: + claimants = sorted(f"{cls.__module__}.{cls.__qualname__}" + for cls in roots) + raise ValueError(f"Unrelated classes claim the name {name!r}: " + f"{', '.join(claimants)}") + return roots[0] if roots else None + + class _DummyFile(io.StringIO): """Dummy file object used by nostdout().""" diff --git a/tests/envs/test_base_env.py b/tests/envs/test_base_env.py index f1614ff5c0..075168e96c 100644 --- a/tests/envs/test_base_env.py +++ b/tests/envs/test_base_env.py @@ -9,6 +9,7 @@ import predicators.envs from predicators import utils from predicators.envs import BaseEnv, create_new_env, get_or_create_env +from predicators.envs.cover import CoverEnv from tests.approaches.test_oracle_approach import ENV_NAME_AND_CLS _MODULE_PATH = predicators.envs.__name__ @@ -42,6 +43,25 @@ def test_env_creation(): create_new_env("Not a real env") +def test_env_creation_ignores_subclasses_that_inherit_the_name(): + """A test's instrumented env never stands in for the env it extends.""" + + class _InstrumentedCoverEnv(CoverEnv): + """Inherits the name ``cover``.""" + + utils.reset_config({ + "env": "cover", + "num_train_tasks": 1, + "num_test_tasks": 1 + }) + # The instrumented class first, as set order sometimes has it. + with patch.object(utils, + "get_all_subclasses", + return_value=[_InstrumentedCoverEnv, CoverEnv]): + env = create_new_env("cover", do_cache=False) + assert env.__class__ is CoverEnv + + @pytest.mark.parametrize("env_name", ("cover", "sandwich")) def test_load_task_from_json(env_name): """Tests for env._load_task_from_json().""" diff --git a/tests/test_utils.py b/tests/test_utils.py index 4cacd0c86d..e15784591d 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -3810,3 +3810,61 @@ def _always_fails(obj, *args, **kwargs): assert "_abc_data" in str(excinfo.value), \ "a genuinely unpicklable object must still report why" assert path.stat().st_size == 0, "a failed dump left a partial file" + + +def test_get_registered_subclass(monkeypatch): + """A name resolves to the most general concrete class that claims it.""" + + class _Base(abc.ABC): + + @classmethod + @abc.abstractmethod + def get_name(cls) -> str: + """The registry name.""" + + class _Named(_Base): + + @classmethod + def get_name(cls) -> str: + return "named" + + class _Instrumented(_Named): + """Inherits the name, as a test's instrumented env does.""" + + class _AbstractNamed(_Base): + + @classmethod + def get_name(cls) -> str: + return "abstract_named" + + @abc.abstractmethod + def run(self) -> None: + """Left to subclasses.""" + + class _Concrete(_AbstractNamed): + + def run(self) -> None: + """Nothing to run.""" + + class _First(_Base): + + @classmethod + def get_name(cls) -> str: + return "claimed_twice" + + class _Second(_Base): + + @classmethod + def get_name(cls) -> str: + return "claimed_twice" + + assert utils.get_registered_subclass(_Base, "named") is _Named + assert utils.get_registered_subclass(_Base, "abstract_named") is _Concrete + assert utils.get_registered_subclass(_Base, "unknown") is None + with pytest.raises(ValueError, match="Unrelated classes claim"): + utils.get_registered_subclass(_Base, "claimed_twice") + # Set order differs between processes; the subclass coming first must + # not change the answer. + monkeypatch.setattr(utils, "get_all_subclasses", + lambda _: [_Instrumented, _Named]) + assert utils.get_registered_subclass(_Base, "named") is _Named