[BugFix] MatNet #108

rl4co/models/zoo/matnet/model.py

@@ -5,32 +5,34 @@
 
 from rl4co.models.zoo.pomo.model import POMO
 from rl4co.envs.common.base import RL4COEnvBase
+from rl4co.utils.pylogger import get_pylogger
+
+log = get_pylogger(__name__)
 
 
 class MatNet(POMO):
     def __init__(
         self,
         env: RL4COEnvBase,
         policy: Union[nn.Module, MatNetPolicy] = None,
-        optimizer_kwargs: dict = {"lr": 4 * 1e-4, "weight_decay": 1e-6},
-        lr_scheduler: str = "MultiStepLR",
-        lr_scheduler_kwargs: dict = {"milestones": [2001, 2101], "gamma": 0.1},
-        use_dihedral_8: bool = False,
         num_starts: int = None,
-        train_data_size: int = 10_000,
-        batch_size: int = 200,
         policy_params: dict = {},
-        model_params: dict = {},
         **kwargs,
     ):
         if policy is None:
             policy = MatNetPolicy(env_name=env.name, **policy_params)
 
+        # Check if num_augment is not 0 or if diheral_8 is True
+        if kwargs.get("num_augment", 0) != 0:
+            log.error("MatNet does not use symmetric augmentation. Setting num_augment to 0.")
+        kwargs["num_augment"] = 0
+        if kwargs.get("use_dihedral_8", True):
+            log.error("MatNet does not use symmetric Dihedral Augmentation. Setting use_dihedral_8 to False.")
+        kwargs["use_dihedral_8"] = False
+
         super(MatNet, self).__init__(
             env=env,
             policy=policy,
-            use_dihedral_8=use_dihedral_8,
             num_starts=num_starts,
-            num_augment=0, # NOTE: for MatNet we don't use augmentation
             **kwargs,
-        )
+        )