Position Logger

agents/utils/configuration.py

@@ -28,7 +28,7 @@ def __init__(self,
                  output_path_location: str,
                  n_timesteps: int,
                  local: bool,
-                 episode_len_sec : int,
+                 episode_len_sec : int = 10,
                  waypoint_buffer_size : int=2,
                  k_p : float=1,
                  k_wp : float=1,

aviaries/factories/position_controller_factory_test.py

@@ -0,0 +1,83 @@
+import numpy as np
+from stable_baselines3.common.env_util import make_vec_env
+from aviaries.PositionControllerAviaryTest import PositionControllerAviaryTest
+from gym_pybullet_drones.utils.enums import ObservationType, ActionType
+from trajectories import DiscretizedTrajectory
+from stable_baselines3.common.vec_env import VecEnv
+from agents.utils.configuration import Configuration
+from .base_factory import BaseFactory
+
+
+class PositionControllerFactoryTest(BaseFactory):
+    action_type: ActionType
+    observation_type: ObservationType
+    t_traj: DiscretizedTrajectory
+    n_env_training: int
+    initial_xyzs: np.ndarray
+    seed: int 
+    use_gui_for_test_env: bool
+    output_path_location: str
+
+    def __init__(self,
+                 config: Configuration,
+                 observation_type: ObservationType,
+                 output_folder: str,
+                 use_gui_for_test_env: bool = True,
+                 n_env_training: int=20,
+                 seed: int = 0,
+        ) -> None:
+        super().__init__()
+        initial_xyzs = config.initial_xyzs
+        action_type = config.action_type
+        t_traj = config.t_traj
+
+        self.initial_xyzs = initial_xyzs
+        self.observation_type = observation_type
+        self.action_type = action_type
+        self.t_traj = t_traj
+        self.n_env_training = n_env_training
+        self.seed = seed
+        self.use_gui_for_test_env = use_gui_for_test_env
+
+    def get_train_env(self) -> VecEnv:
+        train_env = make_vec_env(
+            PositionControllerAviaryTest,
+            env_kwargs=dict(
+                target_trajectory=self.t_traj,
+                initial_xyzs=self.initial_xyzs,
+                obs=self.observation_type,
+                act=self.action_type
+            ),
+            n_envs=self.n_env_training,
+            seed=self.seed
+        )
+        return train_env
+
+    def get_eval_env(self):
+        eval_env = PositionControllerAviaryTest(
+            target_trajectory=self.t_traj,
+            initial_xyzs=self.initial_xyzs,
+            obs=self.observation_type,
+            act=self.action_type
+        )
+        return eval_env
+
+    def get_test_env_gui(self):
+        test_env = PositionControllerAviaryTest(
+            target_trajectory=self.t_traj,
+            initial_xyzs=self.initial_xyzs,
+            gui=self.use_gui_for_test_env,
+            obs=self.observation_type,
+            act=self.action_type,
+            record=False
+        )
+        return test_env
+
+    def get_test_env_no_gui(self):
+        test_env_nogui = PositionControllerAviaryTest(
+            target_trajectory=self.t_traj,
+            initial_xyzs=self.initial_xyzs,
+            obs=self.observation_type,
+            act=self.action_type
+        )
+        return test_env_nogui

aviaries/factories/uzh_trajectory_follower_factory.py

@@ -31,12 +31,13 @@ def __init__(self,
                  use_gui_for_test_env: bool = True,
                  n_env_training: int=20,
                  seed: int = 0,
+                 single_traj=False
         ) -> None:
         super().__init__()
         initial_xyzs = config.initial_xyzs
         action_type = config.action_type
         t_traj = config.t_traj
-
+        self.single_traj = single_traj
         self.initial_xyzs = initial_xyzs
         self.observation_type = observation_type
         self.action_type = action_type
@@ -66,7 +67,8 @@ def get_train_env(self) -> VecEnv:
                 k_wp=self.k_wp,
                 k_s=self.k_s,
                 max_reward_distance=self.max_reward_distance,
-                waypoint_dist_tol=self.waypoint_dist_tol
+                waypoint_dist_tol=self.waypoint_dist_tol,
+                one_traj=self.single_traj
             ),
             n_envs=self.n_env_training,
             seed=self.seed
@@ -85,7 +87,8 @@ def get_eval_env(self):
             k_wp=self.k_wp,
             k_s=self.k_s,
             max_reward_distance=self.max_reward_distance,
-            waypoint_dist_tol=self.waypoint_dist_tol
+            waypoint_dist_tol=self.waypoint_dist_tol,
+            one_traj=self.single_traj
         )
         return eval_env
 
@@ -103,7 +106,8 @@ def get_test_env_gui(self):
             k_wp=self.k_wp,
             k_s=self.k_s,
             max_reward_distance=self.max_reward_distance,
-            waypoint_dist_tol=self.waypoint_dist_tol
+            waypoint_dist_tol=self.waypoint_dist_tol,
+            one_traj=self.single_traj
         )
         return test_env
 
@@ -119,6 +123,7 @@ def get_test_env_no_gui(self):
             k_wp=self.k_wp,
             k_s=self.k_s,
             max_reward_distance=self.max_reward_distance,
-            waypoint_dist_tol=self.waypoint_dist_tol
+            waypoint_dist_tol=self.waypoint_dist_tol,
+            one_traj=self.single_traj
         )
         return test_env_nogui

gym_pybullet_drones/envs/BaseAviary.py

@@ -13,6 +13,9 @@
 import pybullet_data
 import gymnasium as gym
 from gym_pybullet_drones.utils.enums import DroneModel, Physics, ImageType
+from typing import Optional
+from gym_pybullet_drones.utils.pos_logger import PositionLogger, PosLoggerConfig
+
 
 
 class BaseAviary(gym.Env):
@@ -36,7 +39,8 @@ def __init__(self,
                  obstacles=False,
                  user_debug_gui=True,
                  vision_attributes=False,
-                 output_folder='results'
+                 output_folder='results',
+                 log_positions=False
                  ):
         """Initialization of a generic aviary environment.
 
@@ -93,6 +97,18 @@ def __init__(self,
         self.USER_DEBUG = user_debug_gui
         self.URDF = self.DRONE_MODEL.value + ".urdf"
         self.OUTPUT_FOLDER = output_folder
+        self.LOG_POSITIONS = log_positions
+        if not log_positions:
+            self.pos_logger = None
+        else:
+            self.LOG_POSITIONS = True
+            config = PosLoggerConfig(
+                max_len=self.PYB_FREQ / self.CTRL_FREQ,
+                log_folder="./pos_logs"
+            )
+            self.pos_logger = PositionLogger(
+                config
+            )
         #### Load the drone properties from the .urdf file #########
         self.M, \
         self.L, \
@@ -345,6 +361,9 @@ def step(self,
             #### Between aggregate steps for certain types of update ###
             if self.PYB_STEPS_PER_CTRL > 1 and self.PHYSICS in [Physics.DYN, Physics.PYB_GND, Physics.PYB_DRAG, Physics.PYB_DW, Physics.PYB_GND_DRAG_DW]:
                 self._updateAndStoreKinematicInformation()
+            if self.LOG_POSITIONS:
+                cur_pos = p.getBasePositionAndOrientation(self.DRONE_IDS[0], physicsClientId=self.CLIENT)[0]
+                self.pos_logger.log_position(cur_pos)
             #### Step the simulation using the desired physics update ##
             for i in range (self.NUM_DRONES):
                 if self.PHYSICS == Physics.PYB:

gym_pybullet_drones/utils/pos_logger.py

@@ -0,0 +1,53 @@
+import numpy as np
+import os 
+import shutil
+class PosLoggerConfig:
+
+    def __init__(self, max_len=1000, log_folder: str="./positions") -> None:
+        self.max_len=max_len
+        self.log_folder = log_folder
+
+class PositionLogger:
+    def __init__(self, config: PosLoggerConfig) -> None:
+        self.max_len = config.max_len
+        self.n_logged_files = 0
+        self.log_folder = config.log_folder
+        self.positions = []
+
+        if os.path.isfile(self.log_folder):
+            os.remove(self.log_folder)
+        if os.path.isdir(self.log_folder):
+            shutil.rmtree(self.log_folder)
+        os.makedirs(self.log_folder, exist_ok=True)
+
+    def log_position(self, pos):
+        pos = np.array(pos)
+        self.positions.append(pos.reshape((1, 3)))
+        if len(self.positions) >= self.max_len:
+            self._log()
+
+    def _log(self):
+        pos_np = np.vstack(self.positions)
+        fname = os.path.join(
+            self.log_folder,
+            f'{self.n_logged_files}.npy'
+        )
+        np.save(
+            fname, pos_np
+        )
+        self.positions = []
+        self.n_logged_files += 1
+
+    def flush(self):
+        self._log()
+
+def load_positions(log_folder: str) -> np.ndarray:
+    all_positions = []
+    filenames = os.listdir(log_folder)
+    filenames = sorted(filenames, key=lambda x: int(x.split('.')[0]))
+    for fname in filenames:
+        fname_full = os.path.join(log_folder, fname)
+        positions = np.load(fname_full)
+        all_positions.append(positions)
+    all_positions = np.concatenate(all_positions)
+    return all_positions

trajectories/trajectory_factory.py

@@ -4,7 +4,7 @@
 from .waypoint import Waypoint
 from typing import List
 import numpy as np 
-from .discretized_trajectory import DiscretizedTrajectory, DiscreteTrajectoryFromContinuous, DiscretizedTrajectoryFromWaypoints
+from .discretized_trajectory import DiscretizedTrajectory, DiscreteTrajectoryFromContinuous, DiscretizedTrajectoryFromWaypoints, DiscretizedTrajFromNumpy
 from typing import Optional
 from .random_trajectories import sample_random_ctrl_points
 from typing import Tuple 
@@ -100,7 +100,12 @@ def get_simple_smooth_trajectory(cls, starting_waypoint: Waypoint, t_waypoints:
     def get_discr_from_wps(cls, t_waypoints: List[Waypoint]) -> DiscretizedTrajectory:
         traj = DiscretizedTrajectoryFromWaypoints(t_waypoints)        
         return traj
-
+
+    @classmethod
+    def get_discr_from_np(cls, fname: str) -> DiscretizedTrajectory:
+        wps_np = np.load(fname)
+        return DiscretizedTrajFromNumpy(wps_np=wps_np)
+
     @classmethod
     def waypoints_from_numpy(cls, t_waypoints: np.ndarray) -> List[Waypoint]:
         res = [