add biblio

drone/drone_env.py

@@ -2,7 +2,7 @@ import numpy as np
 import gymnasium as gym
 from gymnasium.envs.classic_control import utils
 
-from rotation import quat2eul, deg2rad, rad2deg, quatmul, quatinv
+from rotation import *
 from drone_model import Drone
 
 
@@ -14,62 +14,58 @@ class DroneEnv(gym.Env):
         self.tau = 0.02  # seconds between state updates
         
         # Angle at which to fail the episode
-        self.pos_threshold = np.array([2.4, 2.4, 2.4])
-        self.ang_threshold = np.array([15, 15, 15]) * deg2rad
+        self.state = np.zeros(13)
+        self.target = np.zeros(13)
+        self.obs = np.zeros(13)
+        
+        self.pos_threshold = np.array([1., 1., 1.])
+        self.quat_threshold = np.ones(4)
         self.vel_threshold = np.finfo(np.float32).max * np.ones(3)
         self.rot_threshold = np.finfo(np.float32).max * np.ones(3)
 
         high = np.concatenate(
             [
-                self.pos_threshold * 2,
-                self.ang_threshold * 2,
+                self.pos_threshold,
+                self.quat_threshold,
                 self.vel_threshold,
                 self.rot_threshold
             ],
             dtype=np.float32,
         )
+
         self.observation_space = gym.spaces.Box(-high, high, dtype=np.float32)
         self.action_space = gym.spaces.Box(0, 500, shape=(4,), dtype=np.float32)
-        self.state = np.zeros(13)
-        self.target = np.zeros(13)
-        self.obs = np.zeros(12)
+
         self.steps_beyond_terminated = None
 
     def observation(self):
-        pos = self.state[0:3]
-        quat = self.state[3:7]
-        vel = self.state[7:10]
-        rot = self.state[10:13]
-        d_pos = pos - self.target[0:3]
-        d_quat = quatmul(quat, quatinv(self.target[3:7]))
-        d_vel = vel - self.target[7:10]
-        d_rot = rot - self.target[10:13]
-        
-        d_ang = quat2eul(d_quat)
-        obs = np.concatenate([d_pos, d_ang, d_vel, d_rot])
+        d_pos = self.state[0:3] - self.target[0:3]
+        d_quat = quatmul(quatinv(self.target[3:7]), self.state[3:7])
+        d_vel = self.state[7:10] - self.target[7:10]
+        d_rot = self.state[10:13] - self.target[10:13]
+        obs = np.concatenate([d_pos, d_quat, d_vel, d_rot])
         return obs
 
+    def reward(self):
+        dist_pos = np.linalg.norm(self.obs[0:3])
+        dist_quat = quatdist(self.state[3:7], self.target[3:7])
+        terminated = (dist_pos + dist_quat) > 1
+        reward = 1 - (dist_pos + dist_quat) if not terminated else -1  
+        return reward, terminated
+
     def reset(self, seed=None, options=None):
         super().reset(seed=seed)
         # Note that if you use custom reset bounds, it may lead to out-of-bound
         # state/observations.
-        low, high = utils.maybe_parse_reset_bounds(
-            options, 
-            -0.05, # default low
-             0.05  # default high 
-        )  
-        # self.state = self.np_random.uniform(low=low, high=high, size=(13,))
-        self.state = np.array([0, 0, -20, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0])
-        self.steps_beyond_terminated = None
+        low, high = utils.maybe_parse_reset_bounds(options, -0.1, 0.1)  
+        self.state = self.target + np.array([0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
         self.obs = self.observation()
         return self.obs, {}
 
     def step(self, action):
-        new_state = self.drone.step(self.state, action, self.tau)
-        self.state = new_state
+        self.state = self.drone.step(self.state, action, self.tau)
         self.obs = self.observation()
-        terminated = bool((self.obs[0:3] < -self.pos_threshold).any()  or (self.obs[0:3] > self.pos_threshold).any())
-        reward = 1 if not terminated else 0  # Binary sparse rewards
+        reward, terminated = self.reward()
         return self.obs, reward, terminated, False, {}
 
     def render(self, mode='human'):