aaadddiiitttyyyyaaaa/bbbbbbbbbbb

1import numpy as np
2
3class TrafficEnv:
4    def __init__(self):
5        self.max_steps = 50
6        self.reset()
7
8    def reset(self):
9        self.state = np.random.randint(5, 20, size=4)  # N,S,E,W
10        self.emergency_lane = np.random.choice([0,1,2,3])
11        self.steps = 0
12        return self._get_state()
13
14    def _get_state(self):
15        return np.append(self.state, self.emergency_lane)
16
17    def step(self, action):
18        reward = 0
19
20        # -------------------------
21        # ACTION
22        # -------------------------
23        if action == 0:  # NS
24            self.state[0] = max(0, self.state[0] - 5)
25            self.state[1] = max(0, self.state[1] - 5)
26            cleared = [0,1]
27        else:  # EW
28            self.state[2] = max(0, self.state[2] - 5)
29            self.state[3] = max(0, self.state[3] - 5)
30            cleared = [2,3]
31
32        # -------------------------
33        # EMERGENCY
34        # -------------------------
35        if self.emergency_lane in cleared:
36            reward += 50
37            self.emergency_lane = -1
38        else:
39            reward -= 20
40
41        # -------------------------
42        # TRAFFIC FLOW
43        # -------------------------
44        incoming = np.random.randint(0, 5, size=4)
45        self.state += incoming
46
47        # -------------------------
48        # BASE REWARD
49        # -------------------------
50        total_cars = np.sum(self.state)
51        reward -= total_cars
52
53        # -------------------------
54        # DONE
55        # -------------------------
56        self.steps += 1
57        done = self.steps >= self.max_steps
58
59        # -------------------------
60        # INFO (IMPORTANT)
61        # -------------------------
62        info = {
63            "total_cars": int(total_cars),
64            "emergency_lane": int(self.emergency_lane)
65        }
66
67        return self._get_state(), reward, done, info