diff --git a/pyproject.toml b/pyproject.toml
index b85ce0f..5687167 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,3 +1,7 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
 [project]
 name = "2026sel3-project"
 version = "0.1.0"
@@ -41,3 +45,36 @@ dev = [
     "pytest>=8.0.0",
     "ruff>=0.15.2",
 ]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/brittle_star_project", "src/experiment_logger"]
+
+[tool.mypy]
+mypy_path = "src"
+check_untyped_defs = false
+warn_return_any = false
+
+[[tool.mypy.overrides]]
+module = [
+    "jax.*",
+    "flax.*",
+    "wandb.*",
+    "torch.*",
+    "mujoco.*",
+    "mujoco_warp.*",
+    "optax.*",
+    "tyro.*",
+    "biorobot.*",
+    "gymnasium.*",
+    "matplotlib.*",
+    "mediapy.*",
+    "matplotlib.*",
+    "mediapy.*",
+    "pytest.*",
+    "tensorboard.*",
+    "tqdm.*",
+    "numpy.*",
+    "yaml.*",
+    "moojoco.*"
+]
+ignore_missing_imports = true
diff --git a/src/__init__.py b/src/__init__.py
deleted file mode 100644
index 35ecd10..0000000
--- a/src/__init__.py
+++ /dev/null
@@ -1,25 +0,0 @@
-from .brittle_star_project import (
-    ArenaConfig,
-    Backend,
-    BrittleStarEnv,
-    BrittleStarEnvFactory,
-    EnvConfig,
-    MorphologyConfig,
-    Task,
-    simulate_policy,
-    SimulationConfig,
-    ControlPolicy,
-)
-
-__all__ = [
-    "ArenaConfig",
-    "Backend",
-    "BrittleStarEnv",
-    "BrittleStarEnvFactory",
-    "EnvConfig",
-    "MorphologyConfig",
-    "Task",
-    "simulate_policy",
-    "SimulationConfig",
-    "ControlPolicy",
-]
diff --git a/src/MLPs/mlps.py b/src/brittle_star_project/MLPs/mlps.py
similarity index 100%
rename from src/MLPs/mlps.py
rename to src/brittle_star_project/MLPs/mlps.py
diff --git a/src/ppo.py b/src/brittle_star_project/ppo.py
similarity index 100%
rename from src/ppo.py
rename to src/brittle_star_project/ppo.py
diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py
index 1c807d4..bd97b9e 100644
--- a/src/brittle_star_project/trainers/PPOTrainer.py
+++ b/src/brittle_star_project/trainers/PPOTrainer.py
@@ -16,14 +16,14 @@ from experiment_logger import get_logger
 
 from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs
 from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
-from MLPs.mlps import (
+from brittle_star_project.MLPs.mlps import (
     Actor,
     AgentParams,
     GenericDenseLayersWithActivation,
     OneDenseLayerMLP,
     Storage,
 )
-from ppo import PPO
+from brittle_star_project.ppo import PPO
 
 
 @jax.jit
diff --git a/uv.lock b/uv.lock
index 42321e1..01caa5d 100644
--- a/uv.lock
+++ b/uv.lock
@@ -11,7 +11,7 @@ resolution-markers = [
 [[package]]
 name = "2026sel3-project"
 version = "0.1.0"
-source = { virtual = "." }
+source = { editable = "." }
 dependencies = [
     { name = "biorobot" },
     { name = "cleanrl" },