fix: pass wandb usage option to sample_factory

berenslab · Nov 20, 2024 · 4a6249f · 4a6249f
1 parent 6b4bb58
commit 4a6249f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/doom_creator/util/preload.py b/doom_creator/util/preload.py
@@ -3,7 +3,7 @@
 import shutil
 import struct
 from glob import glob
-from typing import Optional
+from typing import Optional, Set, Tuple
 
 from PIL import Image
 from PIL.PngImagePlugin import PngInfo
@@ -105,8 +105,8 @@ def preload_dataset(
             dataset_wrapper.clean(source_dir)
 
 
-def check_preload(cfg: Config, test: bool):
-    needed_types = set()
+def check_preload(cfg: Config, test: bool) -> Tuple[Config, Set[TextureType]]:
+    needed_types: Set[TextureType] = set()
     for type_cfg in cfg.objects.values():
         for actor in type_cfg.actors.values():
             for i in range(len(actor.textures)):

diff --git a/runner/frameworks/rl/sf_framework.py b/runner/frameworks/rl/sf_framework.py
@@ -142,6 +142,7 @@ def to_sf_cfg(cfg: DictConfig) -> Config:
         SFFramework._set_cfg_cli_argument(
             sf_cfg, "train_dir", os.path.join(cfg.path.run_dir, "train_dir")
         )
+        SFFramework._set_cfg_cli_argument(sf_cfg, "with_wandb", cfg.logging.use_wandb)
         SFFramework._set_cfg_cli_argument(sf_cfg, "wandb_dir", cfg.path.wandb_dir)
         return sf_cfg