RL-2109: Make loading the model from input dir optional

(Issue RL-2109 - Load experiment params in rollout)

Author: EnliteAI Bot

maze/conf/conf_rollout.yaml

@@ -42,6 +42,7 @@ seeding:
 # Optionally the environment or/and the wrappers config can be used from the input_dir.
 # Example for using the environment config from the input_dir:
 #use_input_dir_config:
+#  use_input_dir_model: True
 #  use_input_dir_env: True
 #  use_input_dir_wrappers: False
 use_input_dir_config: ~

maze/core/rollout/rollout_runner.py

@@ -90,8 +90,9 @@ def setup(self, cfg: DictConfig) -> None:
                 if 'use_input_dir_wrappers' in cfg['use_input_dir_config'] and cfg['use_input_dir_config']['use_input_dir_wrappers']:
                     cfg.wrappers = input_dir_config.wrappers
 
-                cfg.model = input_dir_config.model
-                cfg.policy.model = input_dir_config.model
+                if 'use_input_dir_model' in cfg['use_input_dir_config'] and cfg['use_input_dir_config']['use_input_dir_model']:
+                    cfg.model = input_dir_config.model
+                    cfg.policy.model = input_dir_config.model
 
             else:
                 raise FileNotFoundError(f'Config file {config_path} not found')

maze/test/core/rollout/test_rollout_from_input_dir.py

@@ -29,14 +29,16 @@ def experiment_out_dir():
 
 @pytest.mark.parametrize("use_input_dir_env", [True, False])
 @pytest.mark.parametrize("use_input_dir_wrappers", [True, False])
-def test_train_and_rollout(experiment_out_dir, use_input_dir_env, use_input_dir_wrappers):
+@pytest.mark.parametrize("use_input_dir_model", [True, False])
+def test_train_and_rollout(experiment_out_dir, use_input_dir_env, use_input_dir_wrappers, use_input_dir_model):
     """Test loading config from the experiment output directory in the rollout run"""
     rollout_hydra_overrides = {
         "runner": "sequential",
         "runner.n_episodes": "2",
         "policy": "torch_policy",
         "+use_input_dir_config.use_input_dir_env": use_input_dir_env,
         "+use_input_dir_config.use_input_dir_wrappers": use_input_dir_wrappers,
+        "+use_input_dir_config.use_input_dir_model": use_input_dir_model,
         "input_dir": experiment_out_dir,
     }