diff --git a/reasoning_gym/arithmetic/__init__.py b/reasoning_gym/arithmetic/__init__.py
index 05d42b6d..a9cbe9bf 100644
--- a/reasoning_gym/arithmetic/__init__.py
+++ b/reasoning_gym/arithmetic/__init__.py
@@ -1,12 +1,22 @@
-from .chain_sum import ChainSum, ChainSumConfig
-from .basic_arithmetic import ArithmeticDataset, ArithmeticDatasetConfig
+"""
+Arithmetic tasks for training reasoning capabilities:
+- Basic arithmetic
+- Chain sums
+- Word problems
+- Leg counting
+"""
+
+from .basic_arithmetic import ArithmeticDataset, ArithmeticDatasetConfig, arithmetic_dataset
+from .chain_sum import ChainSum, ChainSumConfig, chain_sum
 from .leg_counting import LegCountingConfig, LegCountingDataset, leg_counting_dataset
 
 __all__ = [
-    "ChainSum",
-    "ChainSumConfig",
     "ArithmeticDataset",
     "ArithmeticDatasetConfig",
+    "arithmetic_dataset",
+    "ChainSum",
+    "ChainSumConfig",
+    "chain_sum",
     "LegCountingConfig",
     "LegCountingDataset",
     "leg_counting_dataset"