├── drone_sim ├── viz │ ├── logger.py │ ├── __init__.py │ ├── README.md │ ├── visualiser.py │ └── body.py ├── gym │ ├── __init__.py │ ├── README.md │ ├── test.py │ └── simple_drone_env.py └── sim │ ├── __init__.py │ ├── parameters.py │ ├── README.md │ ├── sensors.py │ └── drone.py ├── requirements.txt ├── setup.py ├── .gitignore ├── LICENSE.md ├── test.py └── README.md /drone_sim/viz/logger.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /drone_sim/gym/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | matplotlib 2 | numpy 3 | -------------------------------------------------------------------------------- /drone_sim/gym/README.md: -------------------------------------------------------------------------------- 1 | # DRONE SIM GYM 2 | This folder contains the files for gym environment. -------------------------------------------------------------------------------- /drone_sim/viz/__init__.py: -------------------------------------------------------------------------------- 1 | from drone_sim.viz.body import Body 2 | from drone_sim.viz.visualiser import Graphics -------------------------------------------------------------------------------- /drone_sim/sim/__init__.py: -------------------------------------------------------------------------------- 1 | from drone_sim.sim.drone import Drone 2 | from drone_sim.sim.sensors import Sensor 3 | from drone_sim.sim.sensors import PositionTracker 4 | from drone_sim.sim.sensors import IMU 5 | from drone_sim.sim.parameters import * 6 | -------------------------------------------------------------------------------- /drone_sim/gym/test.py: -------------------------------------------------------------------------------- 1 | import gym 2 | import numpy as np 3 | from simple_drone_env import SimpleDroneEnv 4 | 5 | env = SimpleDroneEnv() 6 | 7 | for t in range(1000): 8 | action = np.array([1.0, 0.99, 0.99, 1.0]) 9 | print(action) 10 | 11 | obs, reward, done, _ = env.step(action) 12 | #print(obs, done) 13 | 14 | print(reward) 15 | 16 | env.render() 17 | 18 | if done: 19 | break -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from os.path import join 3 | import os 4 | 5 | project_name = "drone_sim" 6 | 7 | requirements = open('requirements.txt').read().splitlines() 8 | 9 | setup( 10 | name=project_name, 11 | version="0.1.0", 12 | url="https://github.com/SuhrudhSarathy/drone_sim", 13 | author="Suhrudh Sarathy", 14 | author_email="suhrudhsarathy@gmail.com", 15 | description="Light-weight Simulator for Drones", 16 | license="MIT", 17 | install_requires=requirements, 18 | packages=[ 19 | project_name, 20 | join(project_name, "gym"), 21 | join(project_name, "sim"), 22 | join(project_name, "viz"), 23 | ], 24 | zip_safe=True 25 | ) -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | 5 | # C extensions 6 | *.so 7 | 8 | # Distribution / packaging 9 | bin/ 10 | build/ 11 | develop-eggs/ 12 | dist/ 13 | eggs/ 14 | lib/ 15 | lib64/ 16 | parts/ 17 | sdist/ 18 | var/ 19 | *.egg-info/ 20 | .installed.cfg 21 | *.egg 22 | 23 | # Installer logs 24 | pip-log.txt 25 | pip-delete-this-directory.txt 26 | 27 | # Unit test / coverage reports 28 | .tox/ 29 | .coverage 30 | .cache 31 | nosetests.xml 32 | coverage.xml 33 | 34 | # Translations 35 | *.mo 36 | 37 | # Mr Developer 38 | .mr.developer.cfg 39 | .project 40 | .pydevproject 41 | 42 | # Rope 43 | .ropeproject 44 | 45 | # Django stuff: 46 | *.log 47 | *.pot 48 | 49 | # Sphinx documentation 50 | docs/_build/ -------------------------------------------------------------------------------- /drone_sim/sim/parameters.py: -------------------------------------------------------------------------------- 1 | """File to Store common parameters for the drone""" 2 | 3 | #! -- Simulation parameters --! 4 | DT = 0.004 5 | #! -- ********* --! 6 | 7 | #! -- Drone parameters --! 8 | MASS = 0.468 9 | RADIUS_PROP = 0.1 10 | K = 2.980e-6 11 | B = 1.140e-7 12 | L = 0.225 13 | G = 9.8 14 | 15 | NULL_ROT = 620.2943 16 | 17 | # Inertia 18 | IXX = 4.856e-3 19 | IYY = 4.856e-3 20 | IZZ = 8.801e-3 21 | 22 | # Drag 23 | AX = 0.25 24 | AY = 0.25 25 | AZ = 0.25 26 | 27 | #! -- ********* --! 28 | 29 | #! -- Sensor Parameters --! 30 | # GPS 31 | GPS_MEAN = 0 32 | GPS_STDDEV = 0.1 33 | 34 | # IMU 35 | IMU_MEANS = { 36 | "accelx": 0, 37 | "accely": 0, 38 | "accelz": 0, 39 | 40 | "gyrox": 0, 41 | "gyroy": 0, 42 | "gyroz": 0 43 | } 44 | IMU_STDDEV = { 45 | "accelx": 0.1, 46 | "accely": 0.1, 47 | "accelz": 0.1, 48 | 49 | "gyrox": 0.1, 50 | "gyroy": 0.1, 51 | "gyroz": 0.1 52 | } 53 | #! -- ********* --! 54 | 55 | #! -- Plotting Params --! 56 | PLT_PAUSE = 1e-6 -------------------------------------------------------------------------------- /drone_sim/viz/README.md: -------------------------------------------------------------------------------- 1 | # VIZ 2 | This folder contains all the code required for visualising the Drone Simulation. The files in the folder are: 3 | 1. `body.py`: This implements the visualisation of the body of the drone. The body uses the Rotation matrix from the drone to visualise the attitude of the drone. The body has to be `attached_to` a Drone object in order to visualise it. 4 | 5 | ```python 6 | # This is the drone Object 7 | drone = Drone() 8 | 9 | # This is the body object which is attached to the Drone to plot it 10 | body = Body() 11 | body.attach_to(drone) 12 | ``` 13 | 14 | 2. `visualiser.py`: This implements the `Graphics` or the plot for visualising the drone. The Graphics object plot all the actors added to it. 15 | ```python 16 | graphics = Graphics() 17 | 18 | # This is the drone Object 19 | drone = Drone() 20 | 21 | graphics.add_actor(drone) 22 | ``` 23 | Multi Drone visualiser is also supported. All the drones that are to be visualised should be added to the Graphics object as actors 24 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | Copyright <2022> 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 4 | 5 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 6 | 7 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -------------------------------------------------------------------------------- /drone_sim/sim/README.md: -------------------------------------------------------------------------------- 1 | # SIM 2 | This folder contains python files for running the main math behind the actual Rigid Body Simulations. The files below are 3 | 1. `drone.py`: Simulates a Drone's Rigid Body Dynamics. 4 | 2. `parameters.py`: Contains all the parameters that are being used for simulation. These can be changed by the user, but the effect will be global. 5 | 3. `sensors.py`: Simulate basic sensors such as GPS (`PositionTracker`) and IMU. Noise can be enabled to mimic real world sensor noise. All the sensors inherit the base `Sensor` class. A sensor needs to be `attached to` the drone: 6 | ```python 7 | drone = Drone() 8 | 9 | mysensor = Sensor() # Valid for other classes also (PositionTracker or IMU) 10 | sensor.attach_to(drone) 11 | 12 | data = sensor.sense() # This method can be called to get the data from the sensor 13 | 14 | ``` 15 | 16 | ## References 17 | 1. https://andrew.gibiansky.com/downloads/pdf/Quadcopter%20Dynamics,%20Simulation,%20and%20Control.pdf 18 | 2. https://sal.aalto.fi/publications/pdf-files/eluu11_public.pdf 19 | 20 | ## Future Work 21 | 1. Simulate additional sensors (Altitude Sensors etc.) 22 | 2. Add obstacles and real time collision checking using Ray Casting (Simulate a 3D Lidar) 23 | 3. The simulation runs on simple Euler Integration and none of the ODE's are solved. Maybe solve the ODE's for more robust dynamics. (scipy.Integrate) 24 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | 2 | """Tests a multidrone simulation""" 3 | from drone_sim.sim import Drone 4 | from drone_sim.viz import Body 5 | from drone_sim.viz import Graphics 6 | from drone_sim.sim.parameters import * 7 | 8 | from time import time 9 | 10 | import numpy as np 11 | 12 | def calculate_rotor_vels(f, tphi, ttheta, tpsi): 13 | b = np.array([f, tphi, ttheta, tpsi]).reshape(-1, 1) 14 | 15 | a = np.array( 16 | [ 17 | [K, K, K, K], 18 | [L*K, 0, 0, -L*K], 19 | [0, L*K, -L*K, 0], 20 | [B, -B, B, -B] 21 | ] 22 | ) 23 | 24 | x = np.linalg.solve(a, b) 25 | 26 | return x 27 | 28 | drone = Drone(True) 29 | drone.z = 2.5 30 | 31 | drone.phi = 1 32 | print(drone.phi) 33 | 34 | # Make a body 35 | body = Body() 36 | body.attach_to(drone) 37 | 38 | 39 | # Make Graphics object 40 | ui = Graphics() 41 | ui.add_actor(drone) 42 | 43 | T = (AX * 5)/np.sin(drone.phi) 44 | print(T) 45 | 46 | for i in range(1050): 47 | x = calculate_rotor_vels(T, 0, 0, 0) 48 | # w1, w2, w3, w4 = x[0][0], x[1][0], x[2][0], x[3][0] 49 | 50 | if i < 25: 51 | 52 | w1, w2, w3, w4 = NULL_ROT, NULL_ROT, NULL_ROT, 1.1*NULL_ROT 53 | else: 54 | w1, w2, w3, w4 = NULL_ROT, NULL_ROT, NULL_ROT, NULL_ROT 55 | # print(drone.vx) 56 | 57 | time_ = time() 58 | drone.step([w1, w2, w3, w4]) 59 | ui.update() -------------------------------------------------------------------------------- /drone_sim/viz/visualiser.py: -------------------------------------------------------------------------------- 1 | """This file is for visualising the Drone while it is in flight using Matplotlib""" 2 | 3 | import numpy as np 4 | import matplotlib.pyplot as plt 5 | from mpl_toolkits.mplot3d import Axes3D 6 | 7 | from drone_sim.sim.drone import Drone 8 | from drone_sim.sim.parameters import PLT_PAUSE 9 | 10 | class Graphics: 11 | """This only generates the 3D Plot of the ongoing simulation and has multi drone support""" 12 | def __init__(self): 13 | self.fig = plt.figure(figsize=(12, 8)) 14 | self.ax = self.fig.add_subplot(projection="3d") 15 | 16 | self.actors = [] 17 | 18 | def add_actor(self, drone): 19 | """This add a Drone for the Graphics object to display""" 20 | drone.body.viz_ax = self.ax 21 | self.actors.append(drone) 22 | 23 | def plot_background(self): 24 | self.ax.set_xlim3d([-5.0, 15.0]) 25 | self.ax.set_ylim3d([-5.0, 15.0]) 26 | self.ax.set_zlim3d([-1.0, 5.0]) 27 | 28 | self.ax.set_xlabel("X") 29 | self.ax.set_ylabel("Y") 30 | self.ax.set_zlabel("Z") 31 | 32 | self.ax.set_title("Quadcopter Simulation") 33 | 34 | def plot_actors(self): 35 | plt.cla() 36 | self.plot_background() 37 | for actor in self.actors: 38 | actor.body.plot_body() 39 | actor.body.update_trajectory() 40 | actor.body.plot_trajectory() 41 | plt.pause(PLT_PAUSE) 42 | 43 | def update(self): 44 | self.plot_actors() 45 | 46 | 47 | 48 | 49 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Drone_Sim [WIP] 2 | A simple Drone dynamics simulation written in Python. 3 | 4 | ## Structure 5 | The repository is divided into 3 main sub-modules: 6 | 1. `sim`: The files that simulate the Rigid Body Dynamics, Sensors etc. are here. 7 | 2. `viz`: The files that visualise the Drone Simulation are here. 8 | 3. `env`: A Simple GYM environment is implemented here. 9 | 10 | ## Installation Instructions 11 | 1. Clone the repository: 12 | ```bash 13 | git clone git@github.com:SuhrudhSarathy/drone_sim.git 14 | ``` 15 | 2. Install using pip 16 | ``` 17 | cd drone_sim 18 | pip install . 19 | ``` 20 | 3. Alternatively, you can install it diretly using pip 21 | ```bash 22 | python -m pip install git+ 23 | ``` 24 | - If the installation, doesn't work try updating the pip by running the following command 25 | ```bash 26 | python -m pip install --upgrade pip 27 | ``` 28 | ## Usage 29 | A simple drone can be simulated using 30 | ```python 31 | from drone_sim.sim import Drone 32 | from drone_sim.sim.parameters import NULL_ROT 33 | from drone_sim.viz import Body 34 | from drone_sim.viz import Graphics 35 | from drone_sim.sim import IMU 36 | 37 | drone = Drone(x=0, y=0, z=2.5, enable_death=True) 38 | 39 | body = Body() 40 | imu = IMU() 41 | 42 | body.attach_to(drone) 43 | imu.attach_to(drone) 44 | 45 | ui = Graphics() 46 | ui.add_actor(drone) 47 | 48 | omega = NULL_ROT 49 | 50 | for t in range(1000): 51 | # Steps the simulation by setting Rotor Velocities 52 | drone.step([omega, omega, omega, omega]) 53 | 54 | # Get the sensor data 55 | sensor_data = imu.sense() 56 | 57 | # Update the UI to display the simulation 58 | ui.update() 59 | ``` 60 | -------------------------------------------------------------------------------- /drone_sim/viz/body.py: -------------------------------------------------------------------------------- 1 | """This file contains drawing of the body""" 2 | from drone_sim.sim.parameters import * 3 | from drone_sim.sim.drone import Drone 4 | import numpy as np 5 | 6 | class Body: 7 | def __init__(self, viz_ax=None): 8 | self.drone = None 9 | self.viz_ax = viz_ax 10 | # Store all the endpoints in an array 11 | # We have the drone as with 5 important coordinates 12 | # Left front rotor, Right front Rotor 13 | # Main Body 14 | # Left Rear rotor, Right read Rotor 15 | self.d = L/np.sqrt(2) 16 | self.coords = np.array( 17 | [ 18 | [0, self.d, -self.d, -self.d, self.d], 19 | [0, self.d, self.d, -self.d, -self.d], 20 | [0, 0, 0, 0, 0] 21 | ] 22 | ) 23 | 24 | def attach_to(self, drone): 25 | self.drone = drone 26 | self.drone.attach_body(self) 27 | 28 | self.trajectory = { 29 | "X": [self.drone.x], 30 | "Y": [self.drone.y], 31 | "Z": [self.drone.z] 32 | } 33 | 34 | def plot_body(self): 35 | assert self.drone is not None, "Add the body to a Drone" 36 | assert self.viz_ax is not None, "Don't know where to plot. Pass the axes when constructing the Object" 37 | 38 | # First transform all the points to the global frame 39 | coords = self.drone.R @ self.coords + self.drone.linear_position() 40 | 41 | origin = coords[:, 0] 42 | rf = coords[:, 1] 43 | lf = coords[:, 2] 44 | lr = coords[:, 3] 45 | rr = coords[:, 4] 46 | 47 | self.viz_ax.plot([origin[0], rf[0]], [origin[1], rf[1]], [origin[2], rf[2]], color="red") 48 | self.viz_ax.plot([origin[0], lf[0]], [origin[1], lf[1]], [origin[2], lf[2]], color="blue") 49 | self.viz_ax.plot([origin[0], lr[0]], [origin[1], lr[1]], [origin[2], lr[2]], color="black") 50 | self.viz_ax.plot([origin[0], rr[0]], [origin[1], rr[1]], [origin[2], rr[2]], color="green") 51 | 52 | self.viz_ax.scatter([origin[0]], [origin[1]], origin[2], color="yellow", s=2) 53 | 54 | def update_trajectory(self): 55 | self.trajectory["X"].append(self.drone.x) 56 | self.trajectory["Y"].append(self.drone.y) 57 | self.trajectory["Z"].append(self.drone.z) 58 | 59 | def plot_trajectory(self): 60 | self.viz_ax.plot(self.trajectory["X"], self.trajectory["Y"], self.trajectory["Z"], "gray") 61 | 62 | -------------------------------------------------------------------------------- /drone_sim/sim/sensors.py: -------------------------------------------------------------------------------- 1 | """This file aims to simulate some of the senors for the drone""" 2 | 3 | import numpy as np 4 | from drone_sim.sim.parameters import * 5 | 6 | class Sensor: 7 | def __init__(self): 8 | self.drone = None 9 | 10 | def attach_to(self, drone): 11 | raise NotImplementedError 12 | 13 | def sense(self): 14 | raise NotImplementedError 15 | 16 | def name(self): 17 | return self.__class__.__name__ 18 | 19 | class PositionTracker(Sensor): 20 | def __init__(self, add_noise=True): 21 | super(PositionTracker, self).__init__() 22 | 23 | self.drone = None 24 | 25 | if self.add_noise: 26 | self.mean = GPS_MEAN 27 | self.std_dev = GPS_STDDEV 28 | else: 29 | self.mean = 0 30 | self.std_dev = 0 31 | 32 | def attach_to(self, drone): 33 | self.drone = drone 34 | self.drone.attach_sensor(self) 35 | 36 | def sense(self): 37 | self.x = self.drone.x + (self.mean + np.random.randn() * self.std_dev) 38 | self.y = self.drone.y + (self.mean + np.random.randn() * self.std_dev) 39 | self.z = self.drone.z + (self.mean + np.random.randn() * self.std_dev) 40 | 41 | return [self.x, self.y, self.z] 42 | 43 | class IMU(Sensor): 44 | def __init__(self, add_noise=True): 45 | super(IMU, self).__init__() 46 | 47 | self.drone = None 48 | self.add_noise = add_noise 49 | if self.add_noise: 50 | self.initialise_random_walks() 51 | else: 52 | self.gyrox_bias = 0 53 | self.gyroy_bias = 0 54 | self.gyroz_bias = 0 55 | 56 | def initialise_random_walks(self): 57 | self.gyrox_bias = 0 58 | self.gyroy_bias = 0 59 | self.gyroz_bias = 0 60 | 61 | def update_bias(self): 62 | self.gyrox_bias += (IMU_MEANS["gyrox"] + np.random.randn() * IMU_STDDEV["gyrox"])*DT 63 | self.gyroy_bias += (IMU_MEANS["gyroy"] + np.random.randn() * IMU_STDDEV["gyroy"])*DT 64 | self.gyroz_bias += (IMU_MEANS["gyroz"] + np.random.randn() * IMU_STDDEV["gyroz"])*DT 65 | 66 | def attach_to(self, drone): 67 | self.drone = drone 68 | self.drone.attach_sensor(self) 69 | 70 | def sense(self): 71 | if self.add_noise: 72 | # If noise is added, gyro's drift has a random walk. 73 | wx = self.drone.p + self.gyrox_bias 74 | wy = self.drone.q + self.gyroy_bias 75 | wz = self.drone.r + self.gyroz_bias 76 | 77 | ax = self.drone.acceleration[0] + IMU_MEANS[0] + np.random.randn() * IMU_STDDEV[0] 78 | ay = self.drone.acceleration[1] + IMU_MEANS[1] + np.random.randn() * IMU_STDDEV[1] 79 | az = self.drone.acceleration[2] + IMU_MEANS[2] + np.random.randn() * IMU_STDDEV[2] 80 | 81 | # Update bias 82 | self.update_bias() 83 | 84 | else: 85 | wx = self.drone.p 86 | wy = self.drone.q 87 | wz = self.drone.r 88 | 89 | ax = self.drone.acceleration[0] 90 | ay = self.drone.acceleration[1] 91 | az = self.drone.acceleration[2] 92 | 93 | return [ax, ay, az], [wx, wy, wz] -------------------------------------------------------------------------------- /drone_sim/gym/simple_drone_env.py: -------------------------------------------------------------------------------- 1 | import gym 2 | from gym import spaces 3 | from drone_sim.sim import Drone 4 | from drone_sim.sim import PositionTracker, IMU 5 | from drone_sim.sim.parameters import * 6 | from drone_sim.viz import Graphics 7 | from drone_sim.viz import Body 8 | 9 | import numpy as np 10 | 11 | class SimpleDroneEnv(gym.Env): 12 | metadata = {'render.modes': ["human"]} 13 | 14 | def __init__(self, goal_position=[0, 5, 5]) -> None: 15 | super(SimpleDroneEnv, self).__init__() 16 | 17 | # Initialise the Drone class, attach body and sensors, and make the Graphics object 18 | # Reset will be manually called in the step function 19 | self.drone = Drone(0, 0, 2, enable_death=False) 20 | self.body = Body() 21 | 22 | self.body.attach_to(self.drone) 23 | 24 | self.ui = Graphics() 25 | self.ui.add_actor(self.drone) 26 | 27 | self.goal_position = goal_position 28 | self.init_position = [0, 0, 2] 29 | # Gym 30 | self.obs_low = -np.inf 31 | self.obs_high = -np.inf 32 | 33 | # Define action and observation space 34 | # We will use a continous action space for the values of the motor's rotation 35 | self.action_space = spaces.Box(low=0, high=1, shape=(4,), dtype=np.float32) 36 | self.observation_space = spaces.Box(low=self.obs_low, high=self.obs_high, shape=(3, 3), dtype=np.float32) 37 | self.reward_range = (0, 100) 38 | 39 | # Some rotation values 40 | self.NULL_ROT = NULL_ROT 41 | 42 | def step(self, action): 43 | assert self.action_space.contains(action), f"{action} doesnot exist in the action space" 44 | 45 | # The action is of the type (w1, w2, w3, w4) 46 | self.drone.step(action*NULL_ROT) 47 | 48 | 49 | # Function has to return the observations 50 | observation = np.array( 51 | [ 52 | [self.drone.x, self.drone.y, self.drone.z], 53 | [self.drone.acceleration[0][0], self.drone.acceleration[1][0], self.drone.acceleration[2][0]], 54 | [self.drone.p, self.drone.q, self.drone.r] 55 | ] 56 | ) 57 | 58 | # Reward is calculated as the ratio of 1 - (dist(present, goal)/dist(start, goal)) 59 | dist_to_go = self.dist([self.drone.x, self.drone.y, self.drone.z], [self.init_position[0], self.init_position[1], self.init_position[2]]) 60 | total_dist = self.dist([self.goal_position[0], self.goal_position[1], self.goal_position[2]], [self.init_position[0], self.init_position[1], self.init_position[2]]) 61 | 62 | reward = 1 - dist_to_go/total_dist 63 | 64 | # Termination condition 65 | if abs(self.drone.phi) > np.radians(60.0) or abs(self.drone.theta) > np.radians(60): 66 | done = True 67 | self.drone.__reset__() 68 | reward -= 5 69 | 70 | # Condition 2: If the z altitude goes negative, we reset the simulation 71 | elif self.drone.z < 0: 72 | done = True 73 | self.drone.__reset__() 74 | reward -= 5 75 | 76 | elif dist_to_go < 0.01: 77 | done = True 78 | reward += 10 79 | 80 | else: 81 | done = False 82 | 83 | 84 | return observation, reward, done, {} 85 | 86 | def reset(self): 87 | # Function to reset the simulation 88 | self.drone.__reset__() 89 | observation = np.array( 90 | [ 91 | [self.drone.x, self.drone.y, self.drone.z], 92 | [self.drone.acceleration[0][0], self.drone.acceleration[1][0], self.drone.acceleration[2][0]], 93 | [self.drone.p, self.drone.q, self.drone.r] 94 | ] 95 | ) 96 | 97 | return observation 98 | 99 | def render(self): 100 | # Function to render the simulation 101 | self.ui.update() 102 | 103 | # Helper functions 104 | def dist(self, x1, x2): 105 | x, y, z = x1 106 | X, Y, Z = x2 107 | 108 | return np.sqrt((x-X)**2 + (y-Y)**2 + (z-Z)**2) 109 | -------------------------------------------------------------------------------- /drone_sim/sim/drone.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy import sin as s, cos as c, tan as t 3 | from drone_sim.sim.parameters import * 4 | 5 | class Drone: 6 | def __init__(self, x=0, y=0, z=0.5, enable_death=True): 7 | # Position 8 | self.x, self.y, self.z = x, y, z 9 | 10 | # Roll Pitch Yaw 11 | self.phi, self.theta, self.psi = 0, 0, 0 12 | 13 | # Linear velocities 14 | self.vx, self.vy, self.vz = 0, 0, 0 15 | 16 | # Angular Velocities 17 | self.p, self.q, self.r = 0, 0, 0 18 | 19 | self.linear_position = lambda: np.array([self.x, self.y, self.z]).reshape(3, 1) 20 | self.angular_position = lambda: np.array([self.phi, self.theta, self.psi]).reshape(3, 1) 21 | self.linear_velocity = lambda: np.array([self.vx, self.vy, self.vz]).reshape(3, 1) 22 | self.angular_velocity = lambda: np.array([self.p, self.q, self.r]).reshape(3, 1) 23 | 24 | # Omegas 25 | self.w1 = 0 26 | self.w2 = 0 27 | self.w3 = 0 28 | self.w4 = 0 29 | 30 | # Inertia Matrix 31 | self.inertia = np.diag([IXX, IYY, IZZ]) 32 | 33 | # Drag Matrix 34 | self.drag = np.diag([AX, AY, AZ]) 35 | 36 | # Thrust Vector 37 | self.thrust = np.array( 38 | [ 39 | [0], 40 | [0], 41 | [K*(self.w1**2 + self.w2**2 + self.w3**2 + self.w4**2)] 42 | ]) 43 | 44 | # Torque Vector 45 | self.torque = np.array( 46 | [ 47 | [L*K*(self.w1**2 - self.w3**2)], 48 | [L*K*(self.w2**2 - self.w4**2)], 49 | [B*(self.w1**2 - self.w2**2 + self.w3**2 - self.w4**2)] 50 | ] 51 | ) 52 | # Drag Force Vector 53 | self.fd = -self.drag@self.linear_velocity() 54 | 55 | # Gravity Vector 56 | self.gravity = np.array([0, 0, -G]).reshape(-1, 1) 57 | 58 | # Transformation Matrices 59 | self.R_phi = np.array( 60 | [ 61 | [c(self.phi), -s(self.phi), 0], 62 | [s(self.phi), c(self.phi), 0], 63 | [0, 0, 1] 64 | ] 65 | ) 66 | 67 | self.R_theta = np.array( 68 | [ 69 | [1, 0, 0], 70 | [0, c(self.theta), -s(self.theta)], 71 | [0, s(self.theta), c(self.theta)] 72 | ] 73 | ) 74 | 75 | self.R_psi = np.array( 76 | [ 77 | [c(self.psi), -s(self.psi), 0], 78 | [s(self.psi), c(self.psi), 0], 79 | [0, 0, 1] 80 | ] 81 | ) 82 | 83 | self.R = self.R_phi @ self.R_theta @ self.R_psi 84 | 85 | self.W =np.array( 86 | [ 87 | [1, 0, -s(self.theta)], 88 | [0, c(self.phi), c(self.theta)*s(self.phi)], 89 | [0, -s(self.phi), c(self.theta)*c(self.phi)] 90 | ] 91 | ) 92 | 93 | self.acceleration = np.zeros((3, 1)) 94 | 95 | self.sensors = [] 96 | self.body = None 97 | 98 | # Death 99 | self.enable_death = enable_death 100 | 101 | def __reset__(self): 102 | """Call this function to reset the simulation. This is called in function method""" 103 | # Position 104 | self.x, self.y, self.z = 0, 0, 0.5 105 | 106 | # Roll Pitch Yaw 107 | self.phi, self.theta, self.psi = 0, 0, 0 108 | 109 | # Linear velocities 110 | self.vx, self.vy, self.vz = 0, 0, 0 111 | 112 | # Angular Velocities 113 | self.p, self.q, self.r = 0, 0, 0 114 | 115 | print("LOG: The Drone is dead. Reset Simulation") 116 | 117 | def step(self, velocities): 118 | """Function to step, i.e. set the angular velocties, to be called externally by the user""" 119 | 120 | self.w1, self.w2, self.w3, self.w4 = velocities[0], -velocities[1], velocities[2], -velocities[3] 121 | # Decide on this, whether, you need to update as soon as you step or not 122 | self.update() 123 | 124 | if self.enable_death: 125 | self.death() 126 | 127 | # All State Update functions 128 | def __update_transformations__(self): 129 | self.R_phi = np.array( 130 | [ 131 | [c(self.phi), -s(self.phi), 0], 132 | [s(self.phi), c(self.phi), 0], 133 | [0, 0, 1] 134 | ] 135 | ) 136 | 137 | self.R_theta = np.array( 138 | [ 139 | [1, 0, 0], 140 | [0, c(self.theta), -s(self.theta)], 141 | [0, s(self.theta), c(self.theta)] 142 | ] 143 | ) 144 | 145 | self.R_psi = np.array( 146 | [ 147 | [c(self.psi), -s(self.psi), 0], 148 | [s(self.psi), c(self.psi), 0], 149 | [0, 0, 1] 150 | ] 151 | ) 152 | 153 | self.R = self.R_phi @ self.R_theta @ self.R_psi 154 | 155 | self.W =np.array( 156 | [ 157 | [1, 0, -s(self.theta)], 158 | [0, c(self.phi), c(self.theta)*s(self.phi)], 159 | [0, -s(self.phi), c(self.theta)*c(self.phi)] 160 | ] 161 | ) 162 | 163 | def __update_thrust_and_torque__(self): 164 | self.thrust = np.array( 165 | [ 166 | [0], 167 | [0], 168 | [K*(self.w1**2 + self.w2**2 + self.w3**2 + self.w4**2)] 169 | ]) 170 | 171 | # Torque Vector 172 | self.torque = np.array( 173 | [ 174 | [L*K*(self.w1**2 - self.w3**2)], 175 | [L*K*(self.w2**2 - self.w4**2)], 176 | [B*(self.w1**2 - self.w2**2 + self.w3**2 - self.w4**2)] 177 | ] 178 | ) 179 | 180 | # Drag Force Vector 181 | self.fd = -self.drag @ self.linear_velocity() 182 | 183 | def __update_acceleration__(self): 184 | """Uses the omegas to update acceleration""" 185 | self.acceleration = self.gravity + (1/MASS)*self.R@self.thrust + (1/MASS)*self.fd 186 | 187 | def __update_omega_dot__(self): 188 | """Updates omega_dot to calculate final state vector""" 189 | ang_vel = self.angular_velocity() 190 | cross_pdt = np.cross(ang_vel.reshape(3,), (self.inertia@ang_vel).reshape(3,)).reshape(3, 1) 191 | MM = self.torque - cross_pdt 192 | 193 | w_dot = np.linalg.inv(self.inertia)@MM 194 | 195 | self.p = w_dot[0][0] 196 | self.q = w_dot[1][0] 197 | self.r = w_dot[2][0] 198 | 199 | def update(self): 200 | """This function is called everytime to update the state of the system""" 201 | # At this point, we assume that the angular velocities are set and hence we go on to update 202 | # simulation step. This will finally be updated as a gym environment, hence we can easily call the 203 | # functions defined in the gym environment to update the velocities. 204 | self.__update_transformations__() 205 | self.__update_thrust_and_torque__() 206 | self.__update_acceleration__() 207 | self.__update_omega_dot__() 208 | 209 | angle = self.angular_position() + self.angular_velocity() * DT 210 | 211 | # Set the angles 212 | self.phi = self.normalise_theta(angle[0][0]) 213 | self.theta = self.normalise_theta(angle[1][0]) 214 | self.psi = self.normalise_theta(angle[2][0]) 215 | 216 | vel = self.linear_velocity() + self.acceleration * DT 217 | 218 | # set the velocities 219 | self.vx = vel[0][0] 220 | self.vy = vel[1][0] 221 | self.vz = vel[2][0] 222 | 223 | position = self.linear_position() + self.linear_velocity() * DT 224 | 225 | # set the positions 226 | self.x = position[0][0] 227 | self.y = position[1][0] 228 | self.z = position[2][0] 229 | 230 | #!--- Helper functions ---! 231 | def normalise_theta(self, angle): 232 | """This is used normalise the angle within -pi to pi""" 233 | if angle > np.pi: 234 | while angle > np.pi: 235 | angle -= 2*np.pi 236 | return angle 237 | elif angle < -np.pi: 238 | while angle < np.pi: 239 | angle += 2*np.pi 240 | return angle 241 | return angle 242 | 243 | #!--- Attaching Sensors ---! 244 | def attach_sensor(self, sensor): 245 | """This is called when a sensor is added to the drone""" 246 | self.sensors.append(sensor) 247 | 248 | def list_sensors(self): 249 | """Can be used to list the sensors placed on the drone""" 250 | for sensor in self.sensors: 251 | print(sensor.__class__.__name__) 252 | 253 | #!--- Attach Body ---! 254 | def attach_body(self, body): 255 | """This is called to attach a body to the drone, i.e. use this to visualise the drone""" 256 | self.body = body 257 | 258 | #!--- Death Constraints ---# 259 | def death(self): 260 | """This function is used to terminate the simulation. This can be enabled or disabled in the constuctor 261 | If a death condition is reached, the simulation is reset.""" 262 | 263 | # Condition 1: If the roll or pitch is more than 60 degrees, we reset the simulation 264 | if abs(self.phi) > np.radians(60.0) or abs(self.theta) > np.radians(60): 265 | self.__reset__() 266 | 267 | # Condition 2: If the z altitude goes negative, we reset the simulation 268 | if self.z < 0: 269 | self.__reset__() --------------------------------------------------------------------------------