|
import numpy as np |
|
from dataclasses import dataclass |
|
|
|
import param_ |
|
from drone import Drone |
|
from dronemodel import DroneModel |
|
from settings import Settings |
|
|
|
|
|
@dataclass |
|
class Team: |
|
""" |
|
Creates a team (it is either red or blue / foe or friend |
|
""" |
|
|
|
is_blue: bool |
|
drones: [Drone] |
|
drone_model: DroneModel |
|
weighted_distance: float = 0 |
|
|
|
def reset(self, obs=None): |
|
|
|
self.delta_weighted_distance() |
|
if obs: |
|
for drone, obs in zip(self.drones, obs): |
|
drone.reset(obs=obs) |
|
else: |
|
for drone in self.drones: |
|
drone.reset() |
|
|
|
def get_observation(self) -> np.ndarray: |
|
""" |
|
get the observation for the RL agent |
|
:return: observation in the form of flatten np.arrays of shape(squad_number, 6*squad_size) |
|
""" |
|
obs = np.array([drone.get_observation() for drone in self.drones]) |
|
deads = ~np.array([drone.is_alive for drone in self.drones]) |
|
|
|
return obs, deads |
|
|
|
def step(self, action: np.ndarray): |
|
obs = np.zeros((len(self.drones), 6)) |
|
done = np.zeros((len(self.drones),)) |
|
reward = np.zeros((len(self.drones),)) |
|
infos = [{} for d in range(len(self.drones))] |
|
for index, drone in enumerate(self.drones): |
|
obs[index], reward[index], done[index], infos[index] = drone.step(action[index]) |
|
done = (sum(done) == len(self.drones)) |
|
info = {'oob': 0, 'hits_target': 0, 'ttl': param_.DURATION, 'distance_to_straight_action': 0} |
|
for i in infos: |
|
info['ttl'] = min(info['ttl'], i['ttl']) |
|
info['oob'] += i['oob'] if 'oob' in i else 0 |
|
info['hits_target'] += i['hits_target'] if 'hits_target' in i else 0 |
|
info['delta_distance'] = 0 if self.is_blue else self.delta_weighted_distance() |
|
info['distance_to_straight_action'] += i['distance_to_straight_action'] \ |
|
if 'distance_to_straight_action' in i else 0 |
|
return obs, sum(reward), done, info |
|
|
|
def delta_weighted_distance(self): |
|
|
|
|
|
team_distance = np.array([d.distance() for d in self.drones if d.is_alive]) |
|
weighted_distance = np.sum(np.exp(-0.5 * (team_distance / (Settings.perimeter/2)) ** 2)) |
|
|
|
delta = weighted_distance - self.weighted_distance if 0 < self.weighted_distance else 0 |
|
|
|
self.weighted_distance = weighted_distance |
|
|
|
return delta |
|
|
|
|
|
class BlueTeam(Team): |
|
""" |
|
Creates the blue team |
|
""" |
|
|
|
def __init__(self, number_of_drones: int = Settings.blues): |
|
self.is_blue = True |
|
self.drone_model = DroneModel(self.is_blue) |
|
|
|
|
|
positions = np.zeros((number_of_drones, 3)) |
|
speeds = np.zeros((number_of_drones, 3)) |
|
blue_speed = Settings.blue_speed_init * self.drone_model.max_speed |
|
circle = index = 0 |
|
for d in range(number_of_drones): |
|
positions[d] = np.array([Settings.blue_circles_rho[circle], |
|
Settings.blue_circles_theta[circle] + index * 2 * np.pi / 3, |
|
Settings.blue_circles_zed[circle]]) |
|
clockwise = 1 - 2 * (circle % 2) |
|
speeds[d] = np.array([blue_speed, np.pi / 6 * clockwise, 0]) |
|
index += 1 |
|
if index == Settings.blues_per_circle[circle]: |
|
index = 0 |
|
circle += 1 |
|
|
|
self.drones = [Drone(is_blue=True, position=position, speed=speed, id_=id_) |
|
for (id_, position, speed) in zip(range(number_of_drones), positions, speeds)] |
|
|
|
|
|
class RedTeam(Team): |
|
""" |
|
Creates the red team |
|
""" |
|
|
|
def __init__(self, number_of_drones: int = Settings.reds): |
|
self.is_blue = False |
|
self.drone_model = DroneModel(self.is_blue) |
|
|
|
positions = np.zeros((number_of_drones, 3)) |
|
positions_noise = np.zeros((number_of_drones, 3)) |
|
speeds = np.zeros((number_of_drones, 3)) |
|
speed_rho = Settings.red_speed_init * self.drone_model.max_speed |
|
squad = index = 0 |
|
for d in range(number_of_drones): |
|
positions[d] = [Settings.red_squads_rho[squad], |
|
Settings.red_squads_theta[squad], |
|
Settings.red_squads_zed[squad]] |
|
positions_noise[d] = [Settings.red_rho_noise[squad], |
|
Settings.red_theta_noise[squad], |
|
Settings.red_zed_noise[squad]] |
|
speeds[d] = [speed_rho, np.pi + positions[d][1], 0] |
|
speeds[d] = [speed_rho, np.pi + positions[d][1], 0] |
|
index += 1 |
|
if index == Settings.red_squads[squad]: |
|
index = 0 |
|
squad += 1 |
|
|
|
self.drones = [Drone(is_blue=False, position=position, position_noise=position_noise, speed=speed, id_=id_) |
|
for (id_, position, position_noise, speed) in |
|
zip(range(len(positions)), positions, positions_noise, speeds)] |
|
|
