Multidrone refactoring (v1)

utiasDSL · Aug 25, 2020 · c1c6697 · c1c6697
1 parent 5899dda
commit c1c6697
Show file tree

Hide file tree

Showing 42 changed files with 1,390 additions and 1,370 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,3 +1,6 @@
+# TO-DO list
+TODOs.md
+
 # macOS users
 .DS_Store
 

diff --git a/LICENSE b/LICENSE
@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2020 JacopoPan
+Copyright (c) 2020 Jacopo Panerati
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

diff --git a/README.md b/README.md
diff --git a/examples/compare.py b/examples/compare.py
@@ -0,0 +1,84 @@
+import os
+import time
+from datetime import datetime
+import pdb
+import math
+import numpy as np
+import pybullet as p
+import pickle
+import matplotlib.pyplot as plt
+
+from utils import *
+from gym_pybullet_drones.envs.Aviary import DroneModel, Physics, Aviary
+from gym_pybullet_drones.envs.Logger import Logger
+from gym_pybullet_drones.envs.Control import ControlType, Control
+
+GUI = False
+RECORD_VIDEO = False
+TRACE_FILE = "example_trace.pkl"
+PHYSICS = Physics.PYB
+
+if __name__ == "__main__":
+
+    #### Load a trace and control reference from a .pkl file ###########################################
+    with open(os.path.dirname(os.path.abspath(__file__))+"/../files/"+TRACE_FILE, 'rb') as in_file:
+        TRACE_TIMESTAMPS, TRACE_DATA, TRACE_CTRL_REFERENCE, _, _, _ = pickle.load(in_file)
+
+
+    #TRACE_CTRL_REFERENCE = np.array(TRACE_CTRL_REFERENCE); TRACE_DATA = np.array(TRACE_DATA)
+
+
+
+    #### Compute trace's parameters ####################################################################
+    DURATION_SEC = int(TRACE_TIMESTAMPS[-1]); SIMULATION_FREQ_HZ = int(len(TRACE_TIMESTAMPS)/TRACE_TIMESTAMPS[-1])
+
+    #### Initialize the simulation #####################################################################
+    env = Aviary(drone_model=DroneModel.CF2X, initial_xyzs=np.array([0,0,.1]).reshape(1,3), physics=PHYSICS, \
+                    normalized_spaces=False, freq=SIMULATION_FREQ_HZ, gui=GUI, obstacles=False, record=RECORD_VIDEO)
+    INITIAL_STATE = env.reset(); action = np.zeros(4); pos_err = 9999.
+
+    #### Assuming TRACE_FILE starts at position [0,0,0] and the sim starts at [0,0,INITIAL_STATE[2]] ###
+    TRACE_CTRL_REFERENCE[:,2] = INITIAL_STATE[2]
+
+    #### Initialize the logger #########################################################################
+    logger = Logger(duration_sec=DURATION_SEC, simulation_freq_hz=SIMULATION_FREQ_HZ, num_drones=2)
+
+    #### Initialize the controller #####################################################################
+    ctrl = Control(env, control_type=ControlType.PID)
+
+    #### Run the comparison ############################################################################
+    START = time.time()
+    for i in range(DURATION_SEC*env.SIM_FREQ):
+
+        #### Step the simulation ###########################################################################
+        obs, reward, done, info = env.step(action)
+
+        #### Compute the next action using the set points from the trace file ##############################
+        action, pos_err, yaw_err = ctrl.computeControlFromState(control_timestep=env.TIMESTEP, state=obs, \
+                                                                target_pos=TRACE_CTRL_REFERENCE[i,0:3], target_vel=TRACE_CTRL_REFERENCE[i,3:6])
+
+
+        #### Re-arrange the trace for consistency with the logger  #########################################
+        trace_obs = np.hstack([TRACE_DATA[i,0:3], np.zeros(4), TRACE_DATA[i,6:9], TRACE_DATA[i,3:6], TRACE_DATA[i,9:12], TRACE_DATA[i,12:16]])
+
+        #### Log the trace ################################################################################# 
+        logger.log(drone=0, timestamp=TRACE_TIMESTAMPS[i], state=trace_obs, control=np.hstack([TRACE_CTRL_REFERENCE[i,:], np.zeros(6)]))
+
+        #### Log the simulation ############################################################################ 
+        logger.log(drone=1, timestamp=i/env.SIM_FREQ, state=obs, control=np.hstack([TRACE_CTRL_REFERENCE[i,:], np.zeros(6)]))
+
+        #### Printout ######################################################################################
+        env.render()
+
+        #### Sync the simulation ###########################################################################
+        if GUI: sync(i, START, env.TIMESTEP)
+
+    #### Close the environment #########################################################################
+    env.close()
+
+    #### Save the simulation results ###################################################################
+    logger.save()
+
+    #### Plot the simulation results ###################################################################
+    logger.plot(pwm=True)
+
diff --git a/examples/fly.py b/examples/fly.py
@@ -0,0 +1,81 @@
+import os
+import time
+from datetime import datetime
+import pdb
+import math
+import numpy as np
+import pybullet as p
+import pickle
+import matplotlib.pyplot as plt
+
+from utils import *
+from gym_pybullet_drones.envs.Aviary import DroneModel, Physics, Aviary
+from gym_pybullet_drones.envs.Logger import Logger
+from gym_pybullet_drones.envs.Control import ControlType, Control
+
+DRONE = DroneModel.CF2X
+NUM_DRONES = 3
+GUI = True
+PHYSICS = Physics.PYB
+RECORD_VIDEO = False
+SIMULATION_FREQ_HZ = 240
+CONTROL_FREQ_HZ = 48
+DURATION_SEC = 10
+
+if __name__ == "__main__":
+
+    #### Initialize the simulation #####################################################################
+    H = .5; R = .5; INIT_XYZS = np.array([ [R*np.cos((i/NUM_DRONES)*2*np.pi+np.pi/2), R*np.sin((i/NUM_DRONES)*2*np.pi+np.pi/2)-R, H] for i in range(NUM_DRONES) ])
+    env = Aviary(drone_model=DRONE, num_drones=NUM_DRONES, initial_xyzs=INIT_XYZS, physics=PHYSICS, visibility_radius=10, \
+                    normalized_spaces=False, freq=SIMULATION_FREQ_HZ, gui=GUI, obstacles=True, record=RECORD_VIDEO); env.reset()
+
+    #### Initialize a circular trajectory ##############################################################
+    PERIOD = 10; NUM_WP = CONTROL_FREQ_HZ*PERIOD; TARGET_POS = np.zeros((NUM_WP,3))
+    for i in range(NUM_WP): TARGET_POS[i,:] = R*np.cos((i/NUM_WP)*(2*np.pi)+np.pi/2)+INIT_XYZS[0,0], R*np.sin((i/NUM_WP)*(2*np.pi)+np.pi/2)-R+INIT_XYZS[0,1], INIT_XYZS[0,2]  
+    wp_counters = [ int(i*NUM_WP/NUM_DRONES) for i in range(NUM_DRONES) ]
+
+    #### Initialize the logger #########################################################################
+    logger = Logger(duration_sec=DURATION_SEC, simulation_freq_hz=SIMULATION_FREQ_HZ, num_drones=NUM_DRONES)
+
+    #### Initialize the controllers ####################################################################    
+    ctrl = [Control(env, control_type=ControlType.PID) for i in range(NUM_DRONES)]
+
+    #### Run the simulation ############################################################################
+    CTRL_EVERY_N_STEPS= int(np.floor(env.SIM_FREQ/CONTROL_FREQ_HZ))
+    action = { str(i): np.array([0,0,0,0]) for i in range(NUM_DRONES) } if NUM_DRONES>1 else np.array([0,0,0,0])
+    START = time.time()
+    for i in range(DURATION_SEC*env.SIM_FREQ):
+
+        #### Step the simulation ###########################################################################
+        obs, reward, done, info = env.step(action)
+
+        #### Transform 1-drone obs into the Dict format of multiple drones to simplify the code ############
+        if NUM_DRONES==1: obs = {"0": {"state": obs}}
+
+        #### Compute control at the desired frequency @@@@@#################################################       
+        if i%CTRL_EVERY_N_STEPS==0:
+
+            #### Compute control for the current waypoint ######################################################
+            for j in range(NUM_DRONES): action[str(j)], _, _ = ctrl[j].computeControlFromState(control_timestep=CTRL_EVERY_N_STEPS*env.TIMESTEP, \
+                                                                                state=obs[str(j)]["state"], target_pos=TARGET_POS[wp_counters[j],:])
+
+            #### Go to the next waypoint and loop ##############################################################
+            for j in range(NUM_DRONES): wp_counters[j] = wp_counters[j] + 1 if wp_counters[j]<(NUM_WP-1) else 0
+
+        #### Log the simulation ############################################################################
+        for j in range(NUM_DRONES): logger.log(drone=j, timestamp=i/env.SIM_FREQ, state= obs[str(j)]["state"], control=np.hstack([ TARGET_POS[wp_counters[j],:], np.zeros(9) ]))   
+
+        #### Printout ######################################################################################
+        env.render()
+
+        #### Sync the simulation ###########################################################################
+        if GUI: sync(i, START, env.TIMESTEP)
+
+    #### Close the environment #########################################################################
+    env.close()
+
+    #### Save the simulation results ###################################################################
+    logger.save()
+
+    #### Plot the simulation results ###################################################################
+    logger.plot()
diff --git a/examples/sa_learning.py → examples/learn.py b/examples/sa_learning.py → examples/learn.py
@@ -16,40 +16,35 @@
 from ray.rllib.agents import ppo
 
 from utils import *
-from gym_pybullet_drones.envs.SingleDroneEnv import SingleDroneEnv
+from gym_pybullet_drones.envs.Aviary import Aviary
 
 RLLIB = False
 
 if __name__ == "__main__":
 
-    ####################################################################################################
-    #### Check out the environment's spaces ############################################################
-    ####################################################################################################
-    env = gym.make("single-drone-v0")
-    print("Action space:", env.action_space)
-    print("Observation space:", env.observation_space)
+    #### Check the environment's spaces ################################################################
+    env = gym.make("the-aviary-v1")
+    print("[INFO] Action space:", env.action_space)
+    print("[INFO] Observation space:", env.observation_space)
     check_env(env, warn=True, skip_render_check=True) 
 
-    ####################################################################################################
     #### Train the model ###############################################################################
-    ####################################################################################################
     if not RLLIB:
         model = A2C(MlpPolicy, env, verbose=1)
         model.learn(total_timesteps=500000)
     else:
         config = ppo.DEFAULT_CONFIG.copy()
         config["num_workers"] = 0
-        agent = ppo.PPOTrainer(config, env="single-drone-v0")
+        agent = ppo.PPOTrainer(config, env="the-aviary-v1")
         for i in range(100):
             results = agent.train()
-            print("{:d}: episode_reward max {:f} min {:f} mean {:f}".format(i, results["episode_reward_max"], results["episode_reward_min"], results["episode_reward_mean"]))
+            print("[INFO] {:d}: episode_reward max {:f} min {:f} mean {:f}".format(i, \
+                    results["episode_reward_max"], results["episode_reward_min"], results["episode_reward_mean"]))
         policy = agent.get_policy()
         print(policy.model.base_model.summary())
 
-    ####################################################################################################
     #### Show (and record a video of) the model's performance ##########################################
-    ####################################################################################################
-    env = SingleDroneEnv(normalized_spaces=True, gui=True, record=True)
+    env = Aviary(normalized_spaces=True, gui=True, record=True)
     obs = env.reset()
     start = time.time()
     for i in range(10*env.SIM_FREQ):

diff --git a/examples/physics.py b/examples/physics.py
@@ -0,0 +1,77 @@
+import os
+import time
+import pdb
+import math
+import numpy as np
+import pybullet as p
+
+from utils import *
+from gym_pybullet_drones.envs.Aviary import DroneModel, Aviary
+
+DURATION_SEC = 60
+NUM_RESETS = 0
+
+if __name__ == "__main__":
+
+    #### Initialize the simulation #####################################################################
+    env = Aviary(drone_model=DroneModel.CF2X, initial_xyzs=np.array([-.7, -.5, .3]).reshape(1,3), \
+                    initial_rpys=np.array([0, -30*(np.pi/180), 0]).reshape(1,3), gui=True, obstacles=True); env.reset()
+
+    #### Get PyBullet's and drone's ids ################################################################
+    PYB_CLIENT = env.getPyBulletClient(); DRONE_IDS = env.getDroneIds()
+
+    #### Make the drone weightless #####################################################################
+    p.setGravity(0, 0, 0, physicsClientId=PYB_CLIENT)
+
+    #### Run the simulation ############################################################################
+    START = time.time()
+    for i in range(DURATION_SEC*env.SIM_FREQ):
+
+        #### Apply x-axis force to the base ################################################################
+        # p.applyExternalForce(DRONE_IDS[0], linkIndex=-1, forceObj=[1e-4,0.,0,], posObj=[0.,0.,0.], flags=p.WORLD_FRAME, physicsClientId=PYB_CLIENT)
+        # p.applyExternalForce(DRONE_IDS[0], linkIndex=-1, forceObj=[1e-4,0.,0,], posObj=[0.,0.,0.], flags=p.LINK_FRAME, physicsClientId=PYB_CLIENT)
+
+        #### Apply y-axis force to the base ################################################################
+        # p.applyExternalForce(DRONE_IDS[0], linkIndex=-1, forceObj=[0.,1e-4,0.], posObj=[0.,0.,0.], flags=p.WORLD_FRAME, physicsClientId=PYB_CLIENT)
+        # p.applyExternalForce(DRONE_IDS[0], linkIndex=-1, forceObj=[0.,1e-4,0.], posObj=[0.,0.,0.], flags=p.LINK_FRAME, physicsClientId=PYB_CLIENT)
+
+        #### Apply z-axis force to the base ################################################################
+        # p.applyExternalForce(DRONE_IDS[0], linkIndex=-1, forceObj=[0.,0.,1e-4], posObj=[0.,0.,0.], flags=p.WORLD_FRAME, physicsClientId=PYB_CLIENT)
+        # p.applyExternalForce(DRONE_IDS[0], linkIndex=-1, forceObj=[0.,0.,1e-4], posObj=[0.,0.,0.], flags=p.LINK_FRAME, physicsClientId=PYB_CLIENT)
+        #### To propeller 0 ################################################################################
+        # p.applyExternalForce(DRONE_IDS[0], linkIndex=0, forceObj=[0.,0.,1e-4], posObj=[0.,0.,0.], flags=p.WORLD_FRAME, physicsClientId=PYB_CLIENT)
+        p.applyExternalForce(DRONE_IDS[0], linkIndex=0, forceObj=[0.,0.,1e-4], posObj=[0.,0.,0.], flags=p.LINK_FRAME, physicsClientId=PYB_CLIENT)
+        #### To the center of mass #########################################################################
+        # p.applyExternalForce(DRONE_IDS[0], linkIndex=4, forceObj=[0.,0.,1e-4], posObj=[0.,0.,0.], flags=p.WORLD_FRAME, physicsClientId=PYB_CLIENT)
+        # p.applyExternalForce(DRONE_IDS[0], linkIndex=4, forceObj=[0.,0.,1e-4], posObj=[0.,0.,0.], flags=p.LINK_FRAME, physicsClientId=PYB_CLIENT)
+
+        #### Apply x-axis torque to the base ###############################################################
+        # p.applyExternalTorque(DRONE_IDS[0], linkIndex=-1, torqueObj=[5e-6,0.,0.], flags=p.WORLD_FRAME, physicsClientId=PYB_CLIENT)
+        # p.applyExternalTorque(DRONE_IDS[0], linkIndex=-1, torqueObj=[5e-6,0.,0.], flags=p.LINK_FRAME, physicsClientId=PYB_CLIENT)
+
+        #### Apply y-axis torque to the base ###############################################################
+        # p.applyExternalTorque(DRONE_IDS[0], linkIndex=-1, torqueObj=[0.,5e-6,0.], flags=p.WORLD_FRAME, physicsClientId=PYB_CLIENT)
+        # p.applyExternalTorque(DRONE_IDS[0], linkIndex=-1, torqueObj=[0.,5e-6,0.], flags=p.LINK_FRAME, physicsClientId=PYB_CLIENT)
+
+        #### Apply z-axis torque to the base ###############################################################
+        # p.applyExternalTorque(DRONE_IDS[0], linkIndex=-1, torqueObj=[0.,0.,5e-6], flags=p.WORLD_FRAME, physicsClientId=PYB_CLIENT)
+        # p.applyExternalTorque(DRONE_IDS[0], linkIndex=-1, torqueObj=[0.,0.,5e-6], flags=p.LINK_FRAME, physicsClientId=PYB_CLIENT)
+        #### To propeller 0 ################################################################################
+        # p.applyExternalTorque(DRONE_IDS[0], linkIndex=0, torqueObj=[0.,0.,5e-6], flags=p.WORLD_FRAME, physicsClientId=PYB_CLIENT)
+        # p.applyExternalTorque(DRONE_IDS[0], linkIndex=0, torqueObj=[0.,0.,5e-6], flags=p.LINK_FRAME, physicsClientId=PYB_CLIENT)
+        #### To the center of mass #########################################################################
+        # p.applyExternalTorque(DRONE_IDS[0], linkIndex=4, torqueObj=[0.,0.,5e-6], flags=p.WORLD_FRAME, physicsClientId=PYB_CLIENT)
+        p.applyExternalTorque(DRONE_IDS[0], linkIndex=4, torqueObj=[0.,0.,5e-6], flags=p.LINK_FRAME, physicsClientId=PYB_CLIENT)
+
+        #### Draw base frame ###############################################################################
+        env._showDroneFrame(0)
+
+        #### Step, printout the state, and sync the simulation #############################################
+        p.stepSimulation(physicsClientId=PYB_CLIENT); env.render(); sync(i, START, env.TIMESTEP)
+
+        #### Reset #########################################################################################
+        if i>1 and i%((DURATION_SEC/(NUM_RESETS+1))*env.SIM_FREQ)==0: env.reset()
+
+    #### Close the environment #########################################################################
+    env.close()
+