Hugging Face's logo

Spaces:
Hexamind
/
swarms
Sleeping

App Files Community
swarms / swarm_policy.py
YvesP's picture
YvesP
initial load
a162e39
raw history blame contribute delete
No virus
13.5 kB
from dataclasses import dataclass
import numpy as np
from stable_baselines3 import SAC
from os import path
import param_
from drone import Drone
@dataclass
class SwarmPolicy:
blues: int
reds: int
is_blue: bool
model: object = None
count: int = 0
def __post_init__(self):
dir_path = "policies/last" + f"/b{self.blues}r{self.reds}/"
model_path = dir_path + ("blues_last.zip" if self.is_blue else "reds_last.zip")
if path.exists(model_path):
print("model loaded:" + model_path)
self.model = SAC.load(model_path, verbose=1)
# predicts from the model or from a simple centripete model
def predict(self, obs):
self.count += 1
if self.model:
action, _ = self.model.predict(obs)
# verbose = 'prediction from ' + (' blue model' if self.is_blue else ' red model') + ' at ' + str(self.count)
# print(verbose)
return action
else:
if self.is_blue:
return self._improved_attack_predict(obs)
else:
return self._simple_predict(obs)
# the default policy
def _simple_predict(self, obs):
simple_obs = _decentralise(obs[0:self.blues*6] if self.is_blue else obs[self.blues*6:(self.blues+self.reds)*6])
drone = Drone(is_blue=self.is_blue)
action = np.array([])
nb_drones = self.blues if self.is_blue else self.reds
for d in range(nb_drones):
assign_pos_speed(drone, d, simple_obs)
'''
pos_n, speed_n = simple_obs[d*6:d*6+3], simple_obs[d*6+3:d*6+6]
pos = drone.from_norm(pos_n, drone.max_positions, drone.min_positions)
drone.position = pos
speed = drone.from_norm(speed_n, drone.max_speeds, drone.min_speeds)
drone.speed = speed
'''
action_d = drone.simple_red()
action = np.hstack((action, action_d))
action = _centralise(action)
return action
# the default attack policy
def _attack_predict(self, obs):
def assign_targets(friends_obs, foes_obs):
'''
this current version is simplistic: all friends target the first foe :)
:param obs:
:return:
'''
friends_nb = len(friends_obs) // 6
foes_nb = len(foes_obs) // 6
friends_targets = -np.ones(friends_nb, dtype=int)
while -1 in friends_targets:
for foe in range(foes_nb):
foe_pos = _denorm(foes_obs[foe*6:foe*6+3])
foe_pos_z = foe_pos[0] * np.exp(1j * foe_pos[1])
min_distance = np.inf
closest_friend = -1
for friend in range(friends_nb):
if friends_targets[friend] == -1:
friend_pos = _denorm(friends_obs[friend*6:friend*6+3])
friend_pos_z = friend_pos[0] * np.exp(1j * friend_pos[1])
distance = np.abs(foe_pos_z - friend_pos_z) ** 2 + (friend_pos[2] - foe_pos[2]) ** 2
if distance < min_distance:
min_distance = distance
closest_friend = friend
friends_targets[closest_friend] = foe
return friends_targets
# gets the friends and foes obs
blue_obs = _decentralise(obs[0:self.blues * 6])
red_obs = _decentralise(obs[self.blues * 6:(self.blues + self.reds) * 6])
friends_obs = blue_obs if self.is_blue else red_obs
foes_obs = red_obs if self.is_blue else blue_obs
# assign red targets to blues
friends_targets = assign_targets(friends_obs, foes_obs)
friend_drone = Drone(is_blue=self.is_blue)
foe_drone = Drone(is_blue=not self.is_blue)
action = np.array([])
nb_drones = self.blues if self.is_blue else self.reds
for d in range(nb_drones):
# assign denormalised position and speed (in m and m/s) to foe drone
friend_drone = assign_pos_speed(friend_drone, d, friends_obs)
foe_drone_id = friends_targets[d]
foe_drone = assign_pos_speed(foe_drone, foe_drone_id, foes_obs)
target, time_to_target = calculate_target(friend_drone, foe_drone)
action_d = friend_drone.simple_red(target=target, z_margin=0)
action = np.hstack((action, action_d))
action = _centralise(action)
return action
# the improved manual attack policy
def _improved_attack_predict(self, obs):
# TODO: revamp the algo as follows
# start from closest reds, find all blues that are compatible with some margin
# among those blues, choose the blue whose first target is the latest
# until there is no red left
# in case there are blues left overs, restart the process, or converge to zero, or..
# or we decide in advance how many blues we want on the closest and populate several blues againts reds
# at the beginning
# TODO: check that reds are correctly ordered
# TODO : add margin in the params
# TODO : case of the foe is not reachable
# gets the friends and foes obs
blue_obs = _decentralise(obs[0:self.blues * 6])
red_obs = _decentralise(obs[self.blues * 6:(self.blues + self.reds) * 6])
friends_obs = blue_obs if self.is_blue else red_obs
foes_obs = red_obs if self.is_blue else blue_obs
friends_nb = self.blues if self.is_blue else self.reds
foes_nb = self.reds if self.is_blue else self.blues
friend_drones = []
for friend_id in range(friends_nb):
# assign denormalised position and speed (in m and m/s) to foe drone
friend_drone = Drone(is_blue=self.is_blue)
friend_drone = assign_pos_speed(friend_drone, friend_id, friends_obs)
friend_drones.append(friend_drone)
foe_drones = []
for foe_id in range(foes_nb):
# assign denormalised position and speed (in m and m/s) to foe drone
foe_drone = Drone(is_blue=not self.is_blue)
foe_drone = assign_pos_speed(foe_drone, foe_id, foes_obs)
foe_drones.append(foe_drone)
targets = np.zeros((friends_nb, foes_nb, 3))
best_targets = -np.ones((friends_nb, 3))
times_to_target = -np.ones((friends_nb, foes_nb))
calculation_done = -np.ones(friends_nb)
friend_chosen = -np.ones(friends_nb)
foe_id = 0
friends_chosen = 0
while foe_id < foes_nb-1 and friends_chosen < friends_nb:
best_friend = -1
best_target = np.zeros(3)
longest_time = -np.inf
foe_drone = foe_drones[foe_id]
for friend_id in range(friends_nb):
if friend_chosen[friend_id] == -1: # the friend has no foe target assigned
friend_drone = friend_drones[friend_id]
if calculation_done[friend_id] == -1: # it has not already been calculated
target_, time_to_target, is_a_catch = calculate_target(friend_drone, foe_drone)
times_to_target[friend_id][foe_id] = time_to_target if is_a_catch else np.inf
targets[friend_id][foe_id] = target_
if times_to_target[friend_id][foe_id] < np.inf: # it is a catch
if calculation_done[friend_id] == -1: # calculation of time with other drones has not been done
for foe_idx in range(foe_id + 1, foes_nb):
foex_drone = foe_drones[foe_idx]
target_, time_to_target, is_a_catch = calculate_target(friend_drone, foex_drone)
times_to_target[friend_id][foe_idx] = time_to_target if is_a_catch else np.inf
targets[friend_id][foe_idx] = target_
calculation_done[friend_id] = 1
closest_target = np.min(times_to_target[friend_id, foe_id+1:])
if longest_time < closest_target:
longest_time = closest_target
best_friend = friend_id
best_target = targets[friend_id][foe_id]
best_targets[best_friend] = best_target
friend_chosen[best_friend] = foe_id
friends_chosen += 1
foe_id += 1
if friends_chosen < friends_nb:
last_foe = foes_nb - 1
for friend_id in range(friends_nb):
if friend_chosen[friend_id] == -1:
if times_to_target[friend_id, last_foe] == -1:
friend_drone, foe_drone = friend_drones[friend_id], foe_drones[last_foe]
target_, time_to_target, is_a_catch = calculate_target(friend_drone, foe_drone)
targets[friend_id][last_foe] = target_
closest_target_id = np.argmin(times_to_target[friend_id, :])
best_targets[friend_id] = targets[friend_id][closest_target_id]
action = np.array([])
for friend_id in range(friends_nb):
action_d = friend_drones[friend_id].simple_red(target=best_targets[friend_id], z_margin=0)
action = np.hstack((action, action_d))
action = _centralise(action)
return action
def assign_pos_speed(drone: Drone, d: int, obs: np.ndarray) -> Drone:
# assign denormalised position and speed (in m and m/s) to friend drone
d = int(d)
pos_n, speed_n = obs[d*6:d*6+3], obs[d*6+3:d*6+6]
pos = drone.from_norm(pos_n, drone.max_positions, drone.min_positions)
drone.position = pos
speed = drone.from_norm(speed_n, drone.max_speeds, drone.min_speeds)
drone.speed = speed
return drone
def _denorm(pos): # from norm (i.e. already decentralised) to meter
drone = Drone()
pos_meter = drone.from_norm(pos, drone.max_positions, drone.min_positions)
return pos_meter
def _decentralise(obs): # [-1,1] to [0,1]
obs = (obs+1)/2
return obs
def _centralise(act): # [0,1] to [-1,1]
act = (act - 1/2) * 2
return act
def calculate_target(blue_drone: Drone, red_drone: Drone) -> (np.ndarray(3, ), float, bool):
'''
:param blue_drone:
:param red_drone:
:return:
'''
# TODO : be more precise at the end of the discovery process
def transform(pos, delta_, theta_):
pos[0] -= delta_
pos[1] -= theta_
return pos[0] * np.exp(1j * pos[1])
def untransform_to_array(pos, delta_, theta_):
pos[0] += delta_
pos[1] += theta_
return pos
theta = red_drone.position[1]
delta = param_.GROUNDZONE
attack_pos = np.copy(blue_drone.position)
target_pos = np.copy(red_drone.position)
z_blue = transform(attack_pos, delta, theta)
z_red = np.real(transform(target_pos, delta, theta))
v_blue = blue_drone.drone_model.max_speed
v_red = red_drone.drone_model.max_speed
blue_shooting_distance = blue_drone.drone_model.distance_to_neutralisation
blue_time_to_zero = (np.abs(z_blue) - blue_shooting_distance) / v_blue
red_time_to_zero = z_red / v_red
if red_time_to_zero <= param_.STEP or red_time_to_zero < blue_time_to_zero + param_.STEP:
return np.zeros(3), red_time_to_zero, False
else:
max_target = z_red
min_target = 0
while True:
target = (max_target + min_target) / 2
blue_time_to_target = max(0, (np.abs(z_blue - target) - blue_shooting_distance) / v_blue)
red_time_to_target = np.abs(z_red - target) / v_red
if red_time_to_target - param_.STEP < blue_time_to_target <= red_time_to_target:
target = untransform_to_array((target / z_red) * target_pos, delta, theta)
return target, blue_time_to_target, True
if red_time_to_target < blue_time_to_target:
max_target = target
min_target = min_target
else: # blue_ time_to_target <= red_time_to_target -1:
max_target = max_target
min_target = target
def unitary_test(rho_blue: float, theta_blue: float, rho_red: float, theta_red: float):
'''
tests for the calculate target function
:param rho_blue:
:param theta_blue:
:param rho_red:
:param theta_red:
:return:
'''
blue_drone = Drone()
blue_drone.position = np.array([rho_blue, theta_blue, 100])
red_drone = Drone(is_blue=False)
red_drone.position = np.array([rho_red, theta_red, 100])
tg, time = calculate_target(blue_drone, red_drone)
print('rho_blue : ', rho_blue, ' theta_blue : ', theta_blue, ' rho_red : ', rho_red, ' theta_red : ', theta_red,
' tg : ', tg, ' time : ', time)
return tg, time
def test():
'''
test for the calculate trajectory function
:return:
'''
for rho_blue in [1000]:
for theta_blue in np.pi * np.array([-1, 0.75, 0.5, 0.25, 0]):
for rho_red in [1000]:
for theta_red in np.pi * np.array([0, 1/4]):
unitary_test(rho_blue=rho_blue, theta_blue=theta_blue, rho_red=rho_red, theta_red=theta_red)
print('done')