fix: sokoban generator fixture

instadeepai · Mar 14, 2024 · 56975ec · 56975ec
1 parent b492701
commit 56975ec
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/jumanji/environments/routing/sokoban/env_test.py b/jumanji/environments/routing/sokoban/env_test.py
@@ -22,7 +22,7 @@
 from jumanji.environments.routing.sokoban.constants import AGENT, BOX, TARGET, WALL
 from jumanji.environments.routing.sokoban.env import Sokoban
 from jumanji.environments.routing.sokoban.generator import (
-    DeepMindGenerator,
+    HuggingFaceDeepMindGenerator,
     SimpleSolveGenerator,
 )
 from jumanji.environments.routing.sokoban.types import State
@@ -33,10 +33,9 @@
 @pytest.fixture(scope="session")
 def sokoban() -> Sokoban:
     env = Sokoban(
-        generator=DeepMindGenerator(
-            difficulty="unfiltered",
-            split="train",
-            proportion_of_files=0.005,
+        generator=HuggingFaceDeepMindGenerator(
+            "unfiltered-train",
+            proportion_of_files=0.01,
         )
     )
     return env