From a0953116f17d2195b6ed56f0c9d9e3a6dcb4c607 Mon Sep 17 00:00:00 2001 From: Dominik Roth Date: Sun, 17 Jul 2022 00:44:16 +0200 Subject: [PATCH] Fixed Bug: Incorrect elipsoids when chols are batched or only diag --- columbus/env.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/columbus/env.py b/columbus/env.py index 86ec75b..627be07 100644 --- a/columbus/env.py +++ b/columbus/env.py @@ -224,13 +224,15 @@ class ColumbusEnv(gym.Env): 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: + while len(chol.shape) > 2: + chol = chol[0] + if chol.shape != (2, 2): + chol = th.diag_embed(chol) + if len(chol.shape) != 2: chol = chol[0] - cov = chol.T @ chol L, V = th.linalg.eig(cov)