From 4048507dfd034f35d1ed21c4d8cf635c07969928 Mon Sep 17 00:00:00 2001 From: Felipe Date: Mon, 26 Apr 2021 21:30:33 -0300 Subject: [PATCH] implement KDTree for vss and vssma envs --- rc_gym/vss/env_ma/vss_gym_ma.py | 35 ++++++++++++++-------- rc_gym/vss/env_vss/vss_gym.py | 52 +++++++++++++-------------------- test.py | 16 ---------- 3 files changed, 43 insertions(+), 60 deletions(-) delete mode 100644 test.py diff --git a/rc_gym/vss/env_ma/vss_gym_ma.py b/rc_gym/vss/env_ma/vss_gym_ma.py index d634e8fd..efaa1f06 100644 --- 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 index 115d5fe1..623001b4 100644 --- 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 deleted file mode 100644 index 0df0a6f6..00000000 --- a/test.py +++ /dev/null @@ -1,16 +0,0 @@ -import gym -import rc_gym - -# Using VSS 3v3 env -env = gym.make('SSLGoToBallIR-v0') - -env.reset() -# Run for 1 episode and print reward at the end -for i in range(1): - done = False - while not done: - # Step using random actions - action = env.action_space.sample() - next_state, reward, done, _ = env.step(action) - env.render() - print(reward) \ No newline at end of file