File size: 2,909 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 |
import numpy as np
import gym
from drone import Drone
class RotateWrapper(gym.Wrapper):
"""
:param env: (gym.Env) Gym environment that will be wrapped
"""
def __init__(self, env):
# Call the parent constructor, so we can access self.env later
super(RotateWrapper, self).__init__(env)
self.angle = 0
def reset(self):
"""
Reset the environment
"""
obs = self.env.reset()
obs = self.post_obs(obs)
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
"""
action = self.rotate_action(action)
obs, reward, done, info = self.env.step(action)
obs = self.post_obs(obs)
return obs, reward, done, info
def post_obs(self, obs):
self.angle = self.get_angle(obs)
return self.rotate_obs(obs)
def get_angle(self, obs: np.ndarray) -> float:
blue_obs, red_obs, blue_fire, red_fire = obs
sigma = 0
for this_obs in (blue_obs, red_obs):
for d in this_obs:
sigma += d[0] * np.exp(1j * d[1])
angle = np.angle(sigma)
return angle
def rotate_obs(self, obs):
blue_obs, red_obs, blue_fire, red_fire = obs
rotated_blue_obs = []
rotated_red_obs = []
for this_obs, is_blue, rotated_obs in zip((blue_obs, red_obs),
(True, False),
(rotated_blue_obs, rotated_red_obs)):
drone = Drone(is_blue=is_blue)
for d in this_obs:
d_meter = np.zeros(6,)
# get the pos and speed in cylindrical coordinated in meters
d_meter[:3] = drone.from_norm(d[:3], drone.max_positions, drone.min_positions)
d_meter[3:6] = drone.from_norm(d[3:6], drone.max_speeds, drone.min_speeds)
# rotate
d_meter[1] -= self.angle
d_meter[4] -= self.angle
# back to norm
d[:3] = drone.to_norm(d_meter[:3], drone.max_positions, drone.min_positions)
d[3:6] = drone.to_norm(d_meter[3:6], drone.max_speeds, drone.min_speeds)
rotated_obs.append(d)
del drone
return np.array(rotated_blue_obs), np.array(rotated_red_obs), blue_fire, red_fire
def rotate_action(self, action):
blue_action, red_action = action
blue_action = np.array(list(map(lambda x: [x[0], (x[1]+self.angle/2/np.pi) % 1, x[2]], blue_action)))
red_action = np.array(list(map(lambda x: [x[0], (x[1]+self.angle/2/np.pi) % 1, x[2]], red_action)))
action = blue_action, red_action
return action
|