diff --git a/reasoning_gym/code/__init__.py b/reasoning_gym/code/__init__.py
index 19aca9df..9a4a9044 100644
--- a/reasoning_gym/code/__init__.py
+++ b/reasoning_gym/code/__init__.py
@@ -5,3 +5,10 @@ Cognition tasks for training reasoning capabilities:
 - Code Execution
 """
 
+from .family_relationships import FamilyRelationshipsConfig, FamilyRelationshipsDataset
+from .bf import BFConfig, BFDataset
+
+__all__ = [
+    "BFConfig",
+    "BFDataset"
+]
diff --git a/reasoning_gym/cognition/__init__.py b/reasoning_gym/cognition/__init__.py
index f5d43196..e1f01947 100644
--- a/reasoning_gym/cognition/__init__.py
+++ b/reasoning_gym/cognition/__init__.py
@@ -8,10 +8,16 @@ Cognition tasks for training reasoning capabilities:
 
 from .color_cube_rotation import ColorCubeRotationConfig, ColorCubeRotationDataset
 from .number_sequences import NumberSequenceConfig, NumberSequenceDataset
+from .rubiks_cube import RubiksCubeConfig, RubiksCubeDataset
+from .figlet_fonts import FigletFontConfig, FigletFontDataset
 
 __all__ = [
     "NumberSequenceConfig",
     "NumberSequenceDataset",
     "ColorCubeRotationConfig",
     "ColorCubeRotationDataset",
+    "RubiksCubeConfig",
+    "RubiksCubeDataset",
+    "FigletFontConfig",
+    "FigletFontDataset"
 ]
diff --git a/reasoning_gym/graphs/__init__.py b/reasoning_gym/graphs/__init__.py
index 8ede1fe9..56231a46 100644
--- a/reasoning_gym/graphs/__init__.py
+++ b/reasoning_gym/graphs/__init__.py
@@ -1,3 +1,9 @@
 from .family_relationships import FamilyRelationshipsConfig, FamilyRelationshipsDataset
+from .quantum_lock import QuantumLockConfig, QuantumLockDataset
 
-__all__ = ["FamilyRelationshipsDataset", "FamilyRelationshipsConfig"]
+__all__ = [
+    "FamilyRelationshipsConfig",
+    "FamilyRelationshipsDataset", 
+    "QuantumLockConfig",
+    "QuantumLockDataset"
+]