Create env.py

env.py

@@ -0,0 +1,55 @@
+import random
+
+class DeliveryEnv:
+    def __init__(self):
+        self.grid_size = 5
+        self.reset()
+
+    def reset(self):
+        self.agent_pos = [0, 0]
+        self.pickup = [2, 2]
+        self.drop = [4, 4]
+        self.has_item = False
+        self.steps = 0
+        return self.state()
+
+    def state(self):
+        return {
+            "agent_pos": self.agent_pos,
+            "has_item": self.has_item,
+            "pickup": self.pickup,
+            "drop": self.drop
+        }
+
+    def step(self, action):
+        reward = -1  # step penalty
+        done = False
+
+        # movement
+        if action == "up":
+            self.agent_pos[0] = max(0, self.agent_pos[0] - 1)
+        elif action == "down":
+            self.agent_pos[0] = min(self.grid_size - 1, self.agent_pos[0] + 1)
+        elif action == "left":
+            self.agent_pos[1] = max(0, self.agent_pos[1] - 1)
+        elif action == "right":
+            self.agent_pos[1] = min(self.grid_size - 1, self.agent_pos[1] + 1)
+
+        # pickup
+        elif action == "pickup":
+            if self.agent_pos == self.pickup and not self.has_item:
+                self.has_item = True
+                reward += 10
+
+        # drop
+        elif action == "drop":
+            if self.agent_pos == self.drop and self.has_item:
+                reward += 20
+                done = True
+
+        self.steps += 1
+
+        if self.steps > 50:
+            done = True
+
+        return self.state(), reward, done, {}