-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathrod_maneuvering_env.py
157 lines (126 loc) · 6.25 KB
/
rod_maneuvering_env.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
import gym
import pygame
from gym import spaces
from rod import Rod
from polygon import Polygon
class RodManeuveringEnv(gym.Env):
def __init__(self, n=30, alpha=0.1, gamma=0.97, epsilon=0.1, theta=0.01, rod_length=180, start_x=90, start_y=450,
goal_x=510, goal_y=180, goal_angle=260):
self.n = n
self.alpha = alpha
self.gamma = gamma
self.epsilon = epsilon
self.theta = theta
self.env_width = 600
self.env_height = 600
self.screen = pygame.display.set_mode((self.env_width, self.env_height))
self.action_space = spaces.Discrete(6)
self.rod = Rod(rod_length, start_x, start_y)
self.goal_rod = Rod(rod_length, goal_x, goal_y, goal_angle)
self.reset()
self.polygons = [Polygon(95, 83, 146, 32, 169, 103, 116, 153),
Polygon(51, 179, 145, 203, 224, 255, 131, 232),
Polygon(129, 280, 120, 412, 227, 488, 236, 360),
Polygon(287, 99, 354, 111, 377, 173, 313, 161),
Polygon(485, 12, 552, 45, 583, 106, 515, 78),
Polygon(433, 145, 469, 270, 447, 394, 413, 273),
Polygon(409, 431, 414, 490, 456, 527, 452, 468),
Polygon(447, 543, 504, 505, 573, 523, 512, 562)]
self.font = pygame.font.Font(pygame.font.get_default_font(), 14)
self.text = self.font.render("Click here to animate the training. But this will slow down the process", True,
(0, 255, 0))
self.textRect = self.text.get_rect()
self.textRect.center = (self.env_width // 2, self.env_height // 2)
def step(self, action):
"""
action == 0 -> up
action == 1 -> down
action == 2 -> right
action == 3 -> left
action == 4 -> +10 degree
action == 5 -> -10 degree
"""
assert self.action_space.contains(action)
reward = 0
rod_positions = self.rod.get_positions()
if self._check_collision(rod_positions[0], rod_positions[1], rod_positions[2], rod_positions[3]):
print("--reset--")
self.reset()
return self.get_obs(), 0, False, {}
if action == 0:
virtual_positions = self.rod.add_virtual_y(-30)
if not self._check_collision(virtual_positions[0], virtual_positions[1], virtual_positions[2],
virtual_positions[3]):
self.rod.center_y -= 30
elif action == 1:
virtual_positions = self.rod.add_virtual_y(30)
if not self._check_collision(virtual_positions[0], virtual_positions[1], virtual_positions[2],
virtual_positions[3]):
self.rod.center_y += 30
elif action == 2:
virtual_positions = self.rod.add_virtual_x(30)
if not self._check_collision(virtual_positions[0], virtual_positions[1], virtual_positions[2],
virtual_positions[3]):
self.rod.center_x += 30
elif action == 3:
virtual_positions = self.rod.add_virtual_x(-30)
if not self._check_collision(virtual_positions[0], virtual_positions[1], virtual_positions[2],
virtual_positions[3]):
self.rod.center_x -= 30
elif action == 4:
if not self._check_collision(self.rod.add_virtual_angle(10)[0], self.rod.add_virtual_angle(10)[1],
self.rod.add_virtual_angle(10)[2], self.rod.add_virtual_angle(10)[3]):
self.rod.angle += 10
elif action == 5:
if not self._check_collision(self.rod.add_virtual_angle(-10)[0], self.rod.add_virtual_angle(-10)[1],
self.rod.add_virtual_angle(-10)[2], self.rod.add_virtual_angle(-10)[3]):
self.rod.angle -= 10
self.rod.fix_angle()
done = False
if self.rod.is_close_to(self.goal_rod):
done = True
reward = 1
self.reset()
return self.get_obs(), reward, done, {}
def get_obs(self):
encrypted_angle = self.rod.angle // 10
encrypted_x = self.rod.center_x // 30
encrypted_y = self.rod.center_y // 30
return encrypted_x, encrypted_y, encrypted_angle
def _draw_line(self, positions, color):
pygame.draw.rect(self.screen, color, [positions[0], positions[1], 3, 3])
pygame.draw.line(self.screen, color, (positions[0], positions[1]), (positions[2], positions[3]))
def _check_collision(self, start_x, start_y, end_x, end_y):
if start_x < 0 or start_x > self.env_width or \
start_y > self.env_height or start_y < 0 or \
end_x > self.env_width or end_x < 0 or \
end_y > self.env_height or end_y < 0:
return True
else:
for polygon in self.polygons:
if polygon.check_collision_with_line(start_x, start_y, end_x, end_y):
return True
return False
def reset(self):
self.rod.reset_position()
def _drawGrid(self):
for i in range(20):
rect = pygame.Rect(30 * i, 0, 0, 600)
pygame.draw.rect(self.screen, (225, 225, 225), rect, 1)
for i in range(20):
rect = pygame.Rect(0, 30 * i, 600, 0)
pygame.draw.rect(self.screen, (225, 225, 225), rect, 1)
def render(self, mode='human'):
rod_positions = self.rod.get_positions()
goal_positions = self.goal_rod.get_positions()
self.screen.fill((255, 255, 255))
self._drawGrid()
self._draw_line(rod_positions, (255, 0, 0))
self._draw_line(goal_positions, (0, 255, 0))
for polygon in self.polygons:
pygame.draw.polygon(self.screen, (128, 128, 128), polygon.get_coordinates())
pygame.display.update()
def render_load_screen(self):
self.screen.fill((0, 0, 0))
self.screen.blit(self.text, self.textRect)
pygame.display.update()