dodox/Columbus
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Toggeable confidence-ellipse (key_c)

Browse Source
This commit is contained in:
Dominik Moritz Roth 2022-07-17 01:00:55 +02:00
parent a0953116f1
commit 54adebfe9b
1 changed files with 35 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

66
columbus/env.py
View File

@ -43,6 +43,7 @@ class ColumbusEnv(gym.Env):
self.draw_observable = True self.draw_observable = True
self.draw_joystick = True self.draw_joystick = True
self.draw_entities = True self.draw_entities = True
self.draw_confidence_ellipse = True
self.void_barrier = True self.void_barrier = True
self.void_damage = 100 self.void_damage = 100
@ -223,41 +224,42 @@ class ColumbusEnv(gym.Env):
pygame.draw.circle(self.screen, smolcol, (20+int(60*x) + pygame.draw.circle(self.screen, smolcol, (20+int(60*x) +
self.joystick_offset[0], 20+int(60*y)+self.joystick_offset[1]), 20, width=0) self.joystick_offset[0], 20+int(60*y)+self.joystick_offset[1]), 20, width=0)
def _draw_confidence_ellipse(self, chol, seconds=1): def _draw_confidence_ellipse(self, chol, forceDraw=False, seconds=1):
col = (255, 255, 255) if (self.draw_confidence_ellipse or forceDraw) and self.visible:
f = seconds/self.speed_fac col = (255, 255, 255)
f = seconds/self.speed_fac
while len(chol.shape) > 2: while len(chol.shape) > 2:
chol = chol[0] chol = chol[0]
if chol.shape != (2, 2): if chol.shape != (2, 2):
chol = th.diag_embed(chol) chol = th.diag_embed(chol)
if len(chol.shape) != 2: if len(chol.shape) != 2:
chol = chol[0] chol = chol[0]
cov = chol.T @ chol cov = chol.T @ chol
L, V = th.linalg.eig(cov) L, V = th.linalg.eig(cov)
L, V = L.real, V.real L, V = L.real, V.real
w, h = int(abs(L[0].item()*f))+1, int(abs(L[1].item()*f))+1 w, h = int(abs(L[0].item()*f))+1, int(abs(L[1].item()*f))+1
# TODO: Is this correct? We try to solve for teh angle from this: # TODO: Is this correct? We try to solve for teh angle from this:
# R = [[cos, -sin],[sin, cos]] # R = [[cos, -sin],[sin, cos]]
# Via only the -sin term. # Via only the -sin term.
#ang1 = int(math.acos(V[0, 0])/math.pi*360) #ang1 = int(math.acos(V[0, 0])/math.pi*360)
ang2 = int(math.asin(-V[0, 1])/math.pi*360) ang2 = int(math.asin(-V[0, 1])/math.pi*360)
#ang3 = int(math.asin(V[1, 0])/math.pi*360) #ang3 = int(math.asin(V[1, 0])/math.pi*360)
ang = ang2 ang = ang2
# print(cov) # print(cov)
#print(w, h, (ang1, ang2, ang3)) #print(w, h, (ang1, ang2, ang3))
x, y = self.agent.pos x, y = self.agent.pos
x, y = x*self.width, y*self.height x, y = x*self.width, y*self.height
rect = pygame.Rect((x-w/2, y-h/2, w, h)) rect = pygame.Rect((x-w/2, y-h/2, w, h))
shape_surface = pygame.Surface(rect.size, pygame.SRCALPHA) shape_surface = pygame.Surface(rect.size, pygame.SRCALPHA)
pygame.draw.ellipse(shape_surface, col, pygame.draw.ellipse(shape_surface, col,
(0, 0, *rect.size), 1) (0, 0, *rect.size), 1)
rotated_surf = pygame.transform.rotate(shape_surface, ang) rotated_surf = pygame.transform.rotate(shape_surface, ang)
self.screen.blit(rotated_surf, rotated_surf.get_rect( self.screen.blit(rotated_surf, rotated_surf.get_rect(
center=rect.center)) center=rect.center))
def _handle_user_input(self): def _handle_user_input(self):
for event in pygame.event.get(): for event in pygame.event.get():
@ -267,6 +269,8 @@ class ColumbusEnv(gym.Env):
self.keypress_timeout = int(self.fps/5) self.keypress_timeout = int(self.fps/5)
if keys[pygame.K_m]: if keys[pygame.K_m]:
self.draw_entities = not self.draw_entities self.draw_entities = not self.draw_entities
elif keys[pygame.K_c]:
self.draw_confidence_ellipse = not self.draw_confidence_ellipse
elif keys[pygame.K_r]: elif keys[pygame.K_r]:
self.reset() self.reset()
elif keys[pygame.K_p]: elif keys[pygame.K_p]: