implement KDTree for vss and vssma envs

jgocm · Apr 30, 2021 · 4048507 · 4048507
1 parent f745d26
commit 4048507
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 60 deletions.
diff --git a/rc_gym/vss/env_ma/vss_gym_ma.py b/rc_gym/vss/env_ma/vss_gym_ma.py
@@ -6,10 +6,11 @@
 import gym
 import numpy as np
 import torch
-from rc_gym.Entities import Frame, Robot
+from rc_gym.Entities import Frame, Robot, Ball
 from rc_gym.Utils.Utils import OrnsteinUhlenbeckAction
 from rc_gym.vss.env_ma.opponent.model import DDPGActor
 from rc_gym.vss.vss_gym_base import VSSBaseEnv
+from rc_gym.Utils import KDTree
 
 
 class VSSMAEnv(VSSBaseEnv):
@@ -236,22 +237,32 @@ def x(): return random.uniform(-field_half_length + 0.1,
         def y(): return random.uniform(-field_half_width + 0.1,
                                        field_half_width - 0.1)
 
-        def theta(): return random.uniform(-180, 180)
+        def theta(): return random.uniform(0, 360)
 
         pos_frame: Frame = Frame()
 
-        pos_frame.ball.x = x()
-        pos_frame.ball.y = y()
-        pos_frame.ball.v_x = 0.
-        pos_frame.ball.v_y = 0.
+        pos_frame.ball = Ball(x=x(), y=y())
 
-        pos_frame.robots_blue[0] = Robot(x=x(), y=y(), theta=theta())
-        pos_frame.robots_blue[1] = Robot(x=x(), y=y(), theta=theta())
-        pos_frame.robots_blue[2] = Robot(x=x(), y=y(), theta=theta())
+        min_dist = 0.1
 
-        pos_frame.robots_yellow[0] = Robot(x=x(), y=y(), theta=theta())
-        pos_frame.robots_yellow[1] = Robot(x=x(), y=y(), theta=theta())
-        pos_frame.robots_yellow[2] = Robot(x=x(), y=y(), theta=theta())
+        places = KDTree()
+        places.insert((pos_frame.ball.x, pos_frame.ball.y))
+
+        for i in range(self.n_robots_blue):
+            pos = (x(), y())
+            while places.get_nearest(pos)[1] < min_dist:
+                pos = (x(), y())
+
+            places.insert(pos)
+            pos_frame.robots_blue[i] = Robot(x=pos[0], y=pos[1], theta=theta())
+
+        for i in range(self.n_robots_yellow):
+            pos = (x(), y())
+            while places.get_nearest(pos)[1] < min_dist:
+                pos = (x(), y())
+
+            places.insert(pos)
+            pos_frame.robots_yellow[i] = Robot(x=pos[0], y=pos[1], theta=theta())
 
         return pos_frame
 

diff --git a/rc_gym/vss/env_vss/vss_gym.py b/rc_gym/vss/env_vss/vss_gym.py
@@ -5,8 +5,9 @@
 
 import gym
 import numpy as np
-from rc_gym.Entities import Frame, Robot
+from rc_gym.Entities import Frame, Robot, Ball
 from rc_gym.vss.vss_gym_base import VSSBaseEnv
+from rc_gym.Utils import KDTree
 
 
 class VSSEnv(VSSBaseEnv):
@@ -200,45 +201,32 @@ def x(): return random.uniform(-field_half_length + 0.1,
         def y(): return random.uniform(-field_half_width + 0.1,
                                        field_half_width - 0.1)
 
-        def theta(): return random.uniform(-180, 180)
+        def theta(): return random.uniform(0, 360)
 
         pos_frame: Frame = Frame()
 
-        pos_frame.ball.x = x()
-        pos_frame.ball.y = y()
-        pos_frame.ball.v_x = 0.
-        pos_frame.ball.v_y = 0.
+        pos_frame.ball = Ball(x=x(), y=y())
 
-        agents = []
-        for i in range(self.n_robots_blue):
-            pos_frame.robots_blue[i] = Robot(x=x(), y=y(), theta=theta())
-            agents.append(pos_frame.robots_blue[i])
-
-        for i in range(self.n_robots_yellow):
-            pos_frame.robots_yellow[i] = Robot(x=x(), y=y(), theta=theta())
-            agents.append(pos_frame.robots_yellow[i])
-
-        def same_position_ref(obj, ref, radius):
-            if obj.x >= ref.x - radius and obj.x <= ref.x + radius and \
-                    obj.y >= ref.y - radius and obj.y <= ref.y + radius:
-                return True
-            return False
-
-        radius_ball = self.field.ball_radius
-        radius_robot = self.field.rbt_radius
-
-        for i in range(len(agents)):
-            while same_position_ref(agents[i], pos_frame.ball, radius_ball):
-                agents[i] = Robot(x=x(), y=y(), theta=theta())
-            for j in range(i):
-                while same_position_ref(agents[i], agents[j], radius_robot):
-                    agents[i] = Robot(x=x(), y=y(), theta=theta())
+        min_dist = 0.1
 
+        places = KDTree()
+        places.insert((pos_frame.ball.x, pos_frame.ball.y))
+
         for i in range(self.n_robots_blue):
-            pos_frame.robots_blue[i] = agents[i]
+            pos = (x(), y())
+            while places.get_nearest(pos)[1] < min_dist:
+                pos = (x(), y())
+
+            places.insert(pos)
+            pos_frame.robots_blue[i] = Robot(x=pos[0], y=pos[1], theta=theta())
 
         for i in range(self.n_robots_yellow):
-            pos_frame.robots_yellow[i] = agents[i+self.n_robots_blue]
+            pos = (x(), y())
+            while places.get_nearest(pos)[1] < min_dist:
+                pos = (x(), y())
+
+            places.insert(pos)
+            pos_frame.robots_yellow[i] = Robot(x=pos[0], y=pos[1], theta=theta())
 
         return pos_frame
 

diff --git a/test.py b/test.py