diff --git a/ml/__init__.py b/ml/__init__.py
index ebf19aa..0f678ef 100644
--- a/ml/__init__.py
+++ b/ml/__init__.py
@@ -1,4 +1,3 @@
-from .env import JackBotEnv
-from .model import ActorCritic
+from .env import JackBotEnv, CurriculumPhase
 
-__all__ = ["JackBotEnv", "ActorCritic"]
+__all__ = ["JackBotEnv", "CurriculumPhase"]
diff --git a/ml/run_eval_training.py b/ml/run_eval_training.py
index 3e84b25..5d95508 100644
--- a/ml/run_eval_training.py
+++ b/ml/run_eval_training.py
@@ -3,10 +3,10 @@ ml/run_eval_training.py - Benchmark reward system across all curriculum phases.
 """
 import time
 import numpy as np
-from ml.env import JackBotEnv, CurriculumPhase
 import sys
 from pathlib import Path
 sys.path.append(str(Path(__file__).resolve().parent.parent))
+from ml.env import JackBotEnv, CurriculumPhase
 
 def evaluate_kinematics(episode_length: int = 1000):
     env = JackBotEnv(