File size: 4,449 Bytes
a162e39 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
from dataclasses import make_dataclass
import numpy as np
from gym import Wrapper
import pandas as pd
import param_
from settings import Settings
Path = make_dataclass("Path", [('path', list), ('step', int), ('d_index', int), ('color', list)])
class MonitorWrapper(Wrapper):
"""
:param env: (gym.Env) Gym environment that will be wrapped
"""
def __init__(self, env, steps, verbose=True):
# Call the parent constructor, so we can access self.env later
super(MonitorWrapper, self).__init__(env)
self.verbose = verbose
self.blue_data = []
self.red_data = []
self.fire_paths = []
lat, lon = Settings.latlon
self.lat_tg = lat
self.lon_tg = lon
self.steps = steps
self.step_ = 0
self.step_max = 0
def reset(self):
"""
Reset the environment
"""
obs = self.env.reset()
self.blue_data = []
self.red_data = []
self.fire_paths = []
self.step_ = 0
self.step_max = 0
return obs
def step(self, action):
"""
:param action: ([float] or int) Action taken by the agent
:return: (np.ndarray, float, bool, dict) observation, reward, is the episode over?, additional informations
"""
obs, reward, done, info = self.env.step(action)
if self.verbose:
self.monitor_state()
self.step_ += 1
if self.step_ == self.steps:
done = True
return obs, reward, done, info
def monitor_state(self):
env = self.env
lat_tg, lon_tg = self.lat_tg, self.lon_tg
for d_index, drone in enumerate(env.blue_team.drones):
lat, lon, zed = drone.to_lat_lon_zed(lat_tg, lon_tg)
self.blue_data.append([self.step_, True, drone.id_, lat, lon, zed, drone.color])
for d_index, drone in enumerate(env.red_team.drones):
lat, lon, zed = drone.to_lat_lon_zed(lat_tg, lon_tg)
self.red_data.append([self.step_, False, drone.id_, lat, lon, zed, drone.color])
for blue_id, red_id in np.argwhere(0 < env.playground.blues_have_fired_reds):
b_lat, b_lon, b_zed = env.blue_team.drones[blue_id].to_lat_lon_zed(lat_tg, lon_tg)
r_lat, r_lon, r_zed = env.red_team.drones[red_id].to_lat_lon_zed(lat_tg, lon_tg)
self.fire_paths.append(Path(step=self.step_,
path=[[b_lat, b_lon, b_zed], [r_lat, r_lon, r_zed]],
color=param_.GREEN_COLOR,
d_index=blue_id))
for red_id, blue_id in np.argwhere(0 < env.playground.reds_have_fired_blues):
b_lat, b_lon, b_zed = env.blue_team.drones[blue_id].to_lat_lon_zed(lat_tg, lon_tg)
r_lat, r_lon, r_zed = env.red_team.drones[red_id].to_lat_lon_zed(lat_tg, lon_tg)
self.fire_paths.append(Path(step=self.step_,
path=[[b_lat, b_lon, b_zed], [r_lat, r_lon, r_zed]],
color=param_.BLACK_COLOR,
d_index=red_id))
def get_df(self):
fire_df = pd.DataFrame(self.fire_paths)
df_columns = ['step', 'isBlue', 'd_index', 'lat', 'lon', 'zed', 'color']
blue_df = pd.DataFrame(self.blue_data, columns=df_columns)
red_df = pd.DataFrame(self.red_data, columns=df_columns)
blue_path_df = []
red_path_df = []
for d_index in range(self.env.nb_blues):
blue_path_df.append(self._get_path_df(blue_df, d_index, color=param_.BLUE_COLOR))
for d_index in range(self.env.nb_reds):
red_path_df.append(self._get_path_df(red_df, d_index, color=param_.RED_COLOR))
return blue_df, red_df, fire_df, blue_path_df, red_path_df
def _get_path_df(self, drone_df: pd.DataFrame, d_index: int, color: int = param_.BLUE_COLOR) -> pd.DataFrame:
traj_length = param_.TRAJ_LENGTH
path_total = drone_df[['lon', 'lat', 'zed', 'step']][drone_df.d_index == d_index].values.tolist()
path = ([Path(path_total[:step+1], step, d_index, color) if step < traj_length
else Path(path_total[step - traj_length:step+1], step, d_index, color)
for step in range(len(path_total))])
path_df = pd.DataFrame(path, columns=['path', 'step', 'd_index', 'color'])
return path_df
|