diff --git a/reasoning_gym/__init__.py b/reasoning_gym/__init__.py
index cb1ba4e6..49068294 100644
--- a/reasoning_gym/__init__.py
+++ b/reasoning_gym/__init__.py
@@ -2,10 +2,20 @@
 Reasoning Gym - A library of procedural dataset generators for training reasoning models
 """
 
+from . import algebra, algorithmic, arithmetic, cognition, data, games, geometry, graphs, logic
 from .factory import create_dataset, register_dataset
 
 __version__ = "0.1.1"
 __all__ = [
+    "algebra",
+    "algorithmic",
+    "arithmetic",
+    "cognition",
+    "data",
+    "games",
+    "geometry",
+    "graphs",
+    "logic",
     "create_dataset",
     "register_dataset",
 ]
diff --git a/reasoning_gym/arithmetic/__init__.py b/reasoning_gym/arithmetic/__init__.py
index 3e916142..9e1a5bc2 100644
--- a/reasoning_gym/arithmetic/__init__.py
+++ b/reasoning_gym/arithmetic/__init__.py
@@ -7,24 +7,36 @@ Arithmetic tasks for training reasoning capabilities:
 - Time intervals
 """
 
-# from .basic_arithmetic import BasicArithmeticDataset
-# from .calendar_arithmetic import CalendarArithmeticDataset
-from .chain_sum import ChainSumDataset
-# from .fraction_simplification import FractionSimplificationDataset
-# from .gcd import GcdDataset
-# from .lcm import LcmDataset
-# from .leg_counting import LegCountingDataset
-# from .prime_factorization import PrimeFactorizationDataset
-# from .time_intervals import TimeIntervalsDataset
+from .basic_arithmetic import BasicArithmeticDataset, BasicArithmeticDatasetConfig
+from .calendar_arithmetic import CalendarArithmeticConfig, CalendarArithmeticDataset
+from .chain_sum import ChainSum, ChainSumConfig
+from .fraction_simplification import FractionSimplificationConfig, FractionSimplificationDataset
+from .gcd import GCDConfig, GCDDataset
+from .lcm import LCMConfig, LCMDataset
+from .leg_counting import LegCountingConfig, LegCountingDataset
+from .prime_factorization import PrimeFactorizationConfig, PrimeFactorizationDataset
+from .time_intervals import TimeIntervalsConfig, TimeIntervalsDataset
 
 __all__ = [
-    # "BasicArithmeticDataset",
-    # "CalendarArithmeticDataset",
-    "ChainSumDataset",
-    # "FractionSimplificationDataset",
-    # "GcdDataset",
-    # "LcmDataset",
-    # "LegCountingDataset",
-    # "PrimeFactorizationDataset",
-    # "TimeIntervalsDataset",
+    "BasicArithmeticDataset",
+    "BasicArithmeticDatasetConfig",
+    "basic_arithmetic_dataset",
+    "ChainSum",
+    "ChainSumConfig",
+    "CalendarArithmeticConfig",
+    "CalendarArithmeticDataset",
+    "Weekday",
+    "CalendarTask",
+    "FractionSimplificationConfig",
+    "FractionSimplificationDataset",
+    "GCDConfig",
+    "GCDDataset",
+    "LCMConfig",
+    "LCMDataset",
+    "LegCountingConfig",
+    "LegCountingDataset",
+    "PrimeFactorizationConfig",
+    "PrimeFactorizationDataset",
+    "TimeIntervalsConfig",
+    "TimeIntervalsDataset",
 ]