Update environment.py

environment.py

@@ -1,9 +1,9 @@
-from typing import Tuple, Dict, Any, Optional
+from typing import Tuple, Dict, Any, Optional
 from models import Observation, Action, Reward, State
 
 class CodeReviewEnv:
-    def __init__(self):
-        self.task = None
+    def __init__(self, task: str = "easy"):
+        self.task = task
         self.reset()
 
     def set_task(self, task: str):
@@ -33,6 +33,7 @@ class CodeReviewEnv:
             raise RuntimeError(f"Invalid task: {self.task}")
 
         return self._get_observation()
+        
 
     def step(self, action: Action) -> Tuple[Observation, Reward, bool, Dict[str, Any]]:
         if self.done: