From 51a734607d4e0bebc6d205ac25a25b12d4cb4ad1 Mon Sep 17 00:00:00 2001 From: Dominik Roth Date: Sat, 16 Jul 2022 23:25:48 +0200 Subject: [PATCH] Added visual for confidence-ellipse --- columbus/env.py | 40 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/columbus/env.py b/columbus/env.py index 188e8c1..86ec75b 100644 --- a/columbus/env.py +++ b/columbus/env.py @@ -7,6 +7,7 @@ import random as random_dont_use from os import urandom import math from columbus import entities, observables +import torch as th class ColumbusEnv(gym.Env): @@ -47,6 +48,7 @@ class ColumbusEnv(gym.Env): self.paused = False self.keypress_timeout = 0 + self.can_accept_chol = True self.rng = random_dont_use.Random() self._seed(self.env_seed) @@ -221,6 +223,40 @@ class ColumbusEnv(gym.Env): 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) + def _draw_confidence_ellipse(self, chol, seconds=1): + # def draw_ellipse_angle(surface, color, rect, angle, width=0): + col = (255, 255, 255) + f = seconds/self.speed_fac + + if len(chol.shape) == 3: + chol = chol[0] + + cov = chol.T @ chol + + L, V = th.linalg.eig(cov) + L, V = L.real, V.real + 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: + # R = [[cos, -sin],[sin, cos]] + # Via only the -sin term. + #ang1 = int(math.acos(V[0, 0])/math.pi*360) + ang2 = int(math.asin(-V[0, 1])/math.pi*360) + #ang3 = int(math.asin(V[1, 0])/math.pi*360) + ang = ang2 + + # print(cov) + #print(w, h, (ang1, ang2, ang3)) + + x, y = self.agent.pos + x, y = x*self.width, y*self.height + rect = pygame.Rect((x-w/2, y-h/2, w, h)) + shape_surface = pygame.Surface(rect.size, pygame.SRCALPHA) + pygame.draw.ellipse(shape_surface, col, + (0, 0, *rect.size), 1) + rotated_surf = pygame.transform.rotate(shape_surface, ang) + self.screen.blit(rotated_surf, rotated_surf.get_rect( + center=rect.center)) + def _handle_user_input(self): for event in pygame.event.get(): pass @@ -251,7 +287,7 @@ class ColumbusEnv(gym.Env): elif keys[pygame.K_d]: self._disturb_next = (1.0, 0.5) - def render(self, mode='human', dont_show=False): + def render(self, mode='human', dont_show=False, chol=None): self._handle_user_input() self.visible = self.visible or not dont_show self._ensure_surface() @@ -267,6 +303,8 @@ class ColumbusEnv(gym.Env): self.screen.blit(self.surface, (0, 0)) self._draw_observable() self._draw_joystick() + if chol != None: + self._draw_confidence_ellipse(chol) if self.visible: pygame.display.update()