diff --git a/gym_pybullet_drones/control/CTBRControl.py b/gym_pybullet_drones/control/CTBRControl.py new file mode 100644 index 000000000..d1b4fb7af --- /dev/null +++ b/gym_pybullet_drones/control/CTBRControl.py @@ -0,0 +1,250 @@ +import os +import numpy as np +import xml.etree.ElementTree as etxml +import pkg_resources +import socket +import struct + +from transforms3d.quaternions import rotate_vector, qconjugate, mat2quat, qmult +from transforms3d.utils import normalized_vector + +from gym_pybullet_drones.utils.enums import DroneModel + +class CTBRControl(object): + """Base class for control. + + Implements `__init__()`, `reset(), and interface `computeControlFromState()`, + the main method `computeControl()` should be implemented by its subclasses. + + """ + + ################################################################################ + + def __init__(self, + drone_model: DroneModel, + g: float=9.8 + ): + """Common control classes __init__ method. + + Parameters + ---------- + drone_model : DroneModel + The type of drone to control (detailed in an .urdf file in folder `assets`). + g : float, optional + The gravitational acceleration in m/s^2. + + """ + #### Set general use constants ############################# + self.DRONE_MODEL = drone_model + """DroneModel: The type of drone to control.""" + self.GRAVITY = g*self._getURDFParameter('m') + """float: The gravitational force (M*g) acting on each drone.""" + self.KF = self._getURDFParameter('kf') + """float: The coefficient converting RPMs into thrust.""" + self.KM = self._getURDFParameter('km') + """float: The coefficient converting RPMs into torque.""" + + self.reset() + + ################################################################################ + + def reset(self): + """Reset the control classes. + + A general use counter is set to zero. + + """ + self.control_counter = 0 + + ################################################################################ + + def computeControlFromState(self, + control_timestep, + state, + target_pos, + target_rpy=np.zeros(3), + target_vel=np.zeros(3), + target_rpy_rates=np.zeros(3) + ): + """Interface method using `computeControl`. + + It can be used to compute a control action directly from the value of key "state" + in the `obs` returned by a call to BaseAviary.step(). + + Parameters + ---------- + control_timestep : float + The time step at which control is computed. + state : ndarray + (20,)-shaped array of floats containing the current state of the drone. + target_pos : ndarray + (3,1)-shaped array of floats containing the desired position. + target_rpy : ndarray, optional + (3,1)-shaped array of floats containing the desired orientation as roll, pitch, yaw. + target_vel : ndarray, optional + (3,1)-shaped array of floats containing the desired velocity. + target_rpy_rates : ndarray, optional + (3,1)-shaped array of floats containing the desired roll, pitch, and yaw rates. + """ + + return self.computeControl(control_timestep=control_timestep, + cur_pos=state[0:3], + cur_quat=np.array([state[6], state[3], state[4], state[5]]), + cur_vel=state[10:13], + cur_ang_vel=state[13:16], + target_pos=target_pos, + target_rpy=target_rpy, + target_vel=target_vel, + target_rpy_rates=target_rpy_rates + ) + + ################################################################################ + + def computeControl(self, + control_timestep, + cur_pos, + cur_quat, + cur_vel, + cur_ang_vel, + target_pos, + target_rpy=np.zeros(3), + target_vel=np.zeros(3), + target_rpy_rates=np.zeros(3) + ): + """Abstract method to compute the control action for a single drone. + + It must be implemented by each subclass of `BaseControl`. + + Parameters + ---------- + control_timestep : float + The time step at which control is computed. + cur_pos : ndarray + (3,1)-shaped array of floats containing the current position. + cur_quat : ndarray + (4,1)-shaped array of floats containing the current orientation as a quaternion. + cur_vel : ndarray + (3,1)-shaped array of floats containing the current velocity. + cur_ang_vel : ndarray + (3,1)-shaped array of floats containing the current angular velocity. + target_pos : ndarray + (3,1)-shaped array of floats containing the desired position. + target_rpy : ndarray, optional + (3,1)-shaped array of floats containing the desired orientation as roll, pitch, yaw. + target_vel : ndarray, optional + (3,1)-shaped array of floats containing the desired velocity. + target_rpy_rates : ndarray, optional + (3,1)-shaped array of floats containing the desired roll, pitch, and yaw rates. + + """ + assert(cur_pos.shape == (3,)), f"cur_pos {cur_pos.shape}" + assert(cur_quat.shape == (4,)), f"cur_quat {cur_quat.shape}" + assert(cur_vel.shape == (3,)), f"cur_vel {cur_vel.shape}" + assert(cur_ang_vel.shape == (3,)), f"cur_ang_vel {cur_ang_vel.shape}" + assert(target_pos.shape == (3,)), f"target_pos {target_pos.shape}" + assert(target_rpy.shape == (3,)), f"target_rpy {target_rpy.shape}" + assert(target_vel.shape == (3,)), f"target_vel {target_vel.shape}" + assert(target_rpy_rates.shape == (3,)), f"target_rpy_rates {target_rpy_rates.shape}" + + G = np.array([.0, .0, -9.8]) + K_P = np.array([3., 3., 8.]) + K_D = np.array([2.5, 2.5, 5.]) + K_RATES = np.array([5., 5., 1.]) + P = target_pos - cur_pos + D = target_vel - cur_vel + tar_acc = K_P * P + K_D * D - G + norm_thrust = np.dot(tar_acc, rotate_vector([.0, .0, 1.], cur_quat)) + # Calculate target attitude + z_body = normalized_vector(tar_acc) + x_body = normalized_vector(np.cross(np.array([.0, 1., .0]), z_body)) + y_body = normalized_vector(np.cross(z_body, x_body)) + tar_att = mat2quat(np.vstack([x_body, y_body, z_body]).T) + # Calculate body rates + q_error = qmult(qconjugate(cur_quat), tar_att) + body_rates = 2 * K_RATES * q_error[1:] + if q_error[0] < 0: + body_rates = -body_rates + + return norm_thrust, *body_rates + +################################################################################ + + def setPIDCoefficients(self, + p_coeff_pos=None, + i_coeff_pos=None, + d_coeff_pos=None, + p_coeff_att=None, + i_coeff_att=None, + d_coeff_att=None + ): + """Sets the coefficients of a PID controller. + + This method throws an error message and exist is the coefficients + were not initialized (e.g. when the controller is not a PID one). + + Parameters + ---------- + p_coeff_pos : ndarray, optional + (3,1)-shaped array of floats containing the position control proportional coefficients. + i_coeff_pos : ndarray, optional + (3,1)-shaped array of floats containing the position control integral coefficients. + d_coeff_pos : ndarray, optional + (3,1)-shaped array of floats containing the position control derivative coefficients. + p_coeff_att : ndarray, optional + (3,1)-shaped array of floats containing the attitude control proportional coefficients. + i_coeff_att : ndarray, optional + (3,1)-shaped array of floats containing the attitude control integral coefficients. + d_coeff_att : ndarray, optional + (3,1)-shaped array of floats containing the attitude control derivative coefficients. + + """ + ATTR_LIST = ['P_COEFF_FOR', 'I_COEFF_FOR', 'D_COEFF_FOR', 'P_COEFF_TOR', 'I_COEFF_TOR', 'D_COEFF_TOR'] + if not all(hasattr(self, attr) for attr in ATTR_LIST): + print("[ERROR] in BaseControl.setPIDCoefficients(), not all PID coefficients exist as attributes in the instantiated control class.") + exit() + else: + self.P_COEFF_FOR = self.P_COEFF_FOR if p_coeff_pos is None else p_coeff_pos + self.I_COEFF_FOR = self.I_COEFF_FOR if i_coeff_pos is None else i_coeff_pos + self.D_COEFF_FOR = self.D_COEFF_FOR if d_coeff_pos is None else d_coeff_pos + self.P_COEFF_TOR = self.P_COEFF_TOR if p_coeff_att is None else p_coeff_att + self.I_COEFF_TOR = self.I_COEFF_TOR if i_coeff_att is None else i_coeff_att + self.D_COEFF_TOR = self.D_COEFF_TOR if d_coeff_att is None else d_coeff_att + + ################################################################################ + + def _getURDFParameter(self, + parameter_name: str + ): + """Reads a parameter from a drone's URDF file. + + This method is nothing more than a custom XML parser for the .urdf + files in folder `assets/`. + + Parameters + ---------- + parameter_name : str + The name of the parameter to read. + + Returns + ------- + float + The value of the parameter. + + """ + #### Get the XML tree of the drone model to control ######## + URDF = self.DRONE_MODEL.value + ".urdf" + path = pkg_resources.resource_filename('gym_pybullet_drones', 'assets/'+URDF) + URDF_TREE = etxml.parse(path).getroot() + #### Find and return the desired parameter ################# + if parameter_name == 'm': + return float(URDF_TREE[1][0][1].attrib['value']) + elif parameter_name in ['ixx', 'iyy', 'izz']: + return float(URDF_TREE[1][0][2].attrib[parameter_name]) + elif parameter_name in ['arm', 'thrust2weight', 'kf', 'km', 'max_speed_kmh', 'gnd_eff_coeff' 'prop_radius', \ + 'drag_coeff_xy', 'drag_coeff_z', 'dw_coeff_1', 'dw_coeff_2', 'dw_coeff_3']: + return float(URDF_TREE[0].attrib[parameter_name]) + elif parameter_name in ['length', 'radius']: + return float(URDF_TREE[1][2][1][0].attrib[parameter_name]) + elif parameter_name == 'collision_z_offset': + COLLISION_SHAPE_OFFSETS = [float(s) for s in URDF_TREE[1][2][0].attrib['xyz'].split(' ')] + return COLLISION_SHAPE_OFFSETS[2] diff --git a/gym_pybullet_drones/envs/BetaAviary.py b/gym_pybullet_drones/envs/BetaAviary.py new file mode 100644 index 000000000..39c449636 --- /dev/null +++ b/gym_pybullet_drones/envs/BetaAviary.py @@ -0,0 +1,309 @@ +import numpy as np +from gymnasium import spaces +import socket +import struct + +from transforms3d.quaternions import rotate_vector, qconjugate + +from gym_pybullet_drones.envs.BaseAviary import BaseAviary +from gym_pybullet_drones.utils.enums import DroneModel, Physics + +class BetaAviary(BaseAviary): + """Multi-drone environment class for use of BetaFlight controller.""" + + ################################################################################ + + def __init__(self, + drone_model: DroneModel=DroneModel.CF2X, + num_drones: int=1, + neighbourhood_radius: float=np.inf, + initial_xyzs=None, + initial_rpys=None, + physics: Physics=Physics.PYB, + pyb_freq: int = 240, + ctrl_freq: int = 240, + gui=False, + record=False, + obstacles=False, + user_debug_gui=True, + output_folder='results', + udp_ip="127.0.0.1" + ): + """Initialization of an aviary environment for use of BetaFlight controller. + + Parameters + ---------- + drone_model : DroneModel, optional + The desired drone type (detailed in an .urdf file in folder `assets`). + num_drones : int, optional + The desired number of drones in the aviary. + neighbourhood_radius : float, optional + Radius used to compute the drones' adjacency matrix, in meters. + initial_xyzs: ndarray | None, optional + (NUM_DRONES, 3)-shaped array containing the initial XYZ position of the drones. + initial_rpys: ndarray | None, optional + (NUM_DRONES, 3)-shaped array containing the initial orientations of the drones (in radians). + physics : Physics, optional + The desired implementation of PyBullet physics/custom dynamics. + pyb_freq : int, optional + The frequency at which PyBullet steps (a multiple of ctrl_freq). + ctrl_freq : int, optional + The frequency at which the environment steps. + gui : bool, optional + Whether to use PyBullet's GUI. + record : bool, optional + Whether to save a video of the simulation in folder `files/videos/`. + obstacles : bool, optional + Whether to add obstacles to the simulation. + user_debug_gui : bool, optional + Whether to draw the drones' axes and the GUI RPMs sliders. + udp_ip : base ip for betaflight controller emulator + + """ + super().__init__(drone_model=drone_model, + num_drones=num_drones, + neighbourhood_radius=neighbourhood_radius, + initial_xyzs=initial_xyzs, + initial_rpys=initial_rpys, + physics=physics, + pyb_freq=pyb_freq, + ctrl_freq=ctrl_freq, + gui=gui, + record=record, + obstacles=obstacles, + user_debug_gui=user_debug_gui, + output_folder=output_folder + ) + + # Initialize connection to BetaFlight Controller + self.UDP_IP = udp_ip + self.ARM_TIME = 1 + self.TRAJ_TIME = 1.5 + + self.sock = socket.socket(socket.AF_INET, # Internet + socket.SOCK_DGRAM) # UDP + + self.sock_pwm = socket.socket(socket.AF_INET, # Internet + socket.SOCK_DGRAM) # UDP + self.sock_pwm.bind((self.UDP_IP, 9002)) + self.sock_pwm.settimeout(0.0) + + self.last_action = np.zeros((self.NUM_DRONES, 4)) + + ################################################################################ + + def step(self, action, i): + # Action is (thro, roll, pitch, yaw) + t = i/self.CTRL_FREQ + + thro = 1000 # Positive up + yaw = 1500 # Positive CCW + pitch = 1500 # Positive forward in x + roll = 1500 # Positive right/forward in y + + if t > self.TRAJ_TIME: + thro, roll, pitch, yaw = self.ctbr2beta(*action[0]) + + #### RC message to Betaflight ############################## + aux1 = 1000 if t < self.ARM_TIME else 1500 # Arm + rc_packet = struct.pack( + '@dHHHHHHHHHHHHHHHH', + t, # datetime.now().timestamp(), # double timestamp; // in seconds + round(roll), round(pitch), round(thro), round(yaw), # roll, pitch, throttle, yaw + aux1, 1000, 1000, 1000, # aux 1, .. + 1000, 1000, 1000, 1000, + 1000, 1000, 1000, 1000 + ) + # print("rc", struct.unpack('@dHHHHHHHHHHHHHHHH', rc_packet)) + self.sock.sendto(rc_packet, (self.UDP_IP, 9004)) + + #### PWMs message from Betaflight ########################## + try: + data, addr = self.sock_pwm.recvfrom(16) # buffer size is 100 bytes (servo_packet size 16) + except socket.error as msg: + _action = self.last_action + pass + else: + # print("received message: ", data) + _action = np.array(struct.unpack('@ffff', data)).reshape((1,4)) + + obs, reward, terminated, truncated, info = super().step(_action) + self.last_action = _action + + #### State message to Betaflight ########################### + o = obs[0] # p, q, euler, v, w, rpm (all in world frame) + p = o[:3] + q = np.array([o[6], o[3], o[4], o[5]]) # w, x, y, z + v = o[10:13] + w = o[13:16] # world frame + w_body = rotate_vector(w, qconjugate(q)) # local frame + + fdm_packet = struct.pack( + '@dddddddddddddddddd', # t, w, a, q, v, p, pressure + t, # double timestamp in seconds + # minus signs due to ENU to NED conversion + w_body[0], -w_body[1], -w_body[2], # double imu_angular_velocity_rpy[3] + 0, 0, 0, # double imu_linear_acceleration_xyz[3] + 1., .0, .0, 0., # double imu_orientation_quat[4] w, x, y, z + 0, 0, 0, # double velocity_xyz[3] + 0, 0, 0, # double position_xyz[3] + 1.0 # double pressure; + ) + self.sock.sendto(fdm_packet, (self.UDP_IP, 9003)) + + return obs, reward, terminated, truncated, info + + ################################################################################ + + def ctbr2beta(self, thrust, roll, pitch, yaw): + MIN_CHANNEL = 1000 + MAX_CHANNEL = 2000 + MAX_RATE = 360 + MAX_THRUST = 40.9 + mid = (MAX_CHANNEL + MIN_CHANNEL) / 2 + d = (MAX_CHANNEL - MIN_CHANNEL) / 2 + thrust = thrust / MAX_THRUST * d * 2 + MIN_CHANNEL + rates = np.array([roll, pitch, -yaw]) + rates = rates / np.pi * 180 / MAX_RATE * d + mid + thrust = np.clip(thrust, MIN_CHANNEL, MAX_CHANNEL) + rates = np.clip(rates, MIN_CHANNEL, MAX_CHANNEL) + return thrust, *rates + + ################################################################################ + + def _actionSpace(self): + """Returns the action space of the environment. + + Returns + ------- + spaces.Box + An ndarray of shape (NUM_DRONES, 4) for the commanded RPMs. + + """ + #### Action vector ######## P0 P1 P2 P3 + act_lower_bound = np.array([[0., 0., 0., 0.] for i in range(self.NUM_DRONES)]) + act_upper_bound = np.array([[self.MAX_RPM, self.MAX_RPM, self.MAX_RPM, self.MAX_RPM] for i in range(self.NUM_DRONES)]) + return spaces.Box(low=act_lower_bound, high=act_upper_bound, dtype=np.float32) + + ################################################################################ + + def _observationSpace(self): + """Returns the observation space of the environment. + + Returns + ------- + spaces.Box + The observation space, i.e., an ndarray of shape (NUM_DRONES, 20). + + """ + #### Observation vector ### X Y Z Q1 Q2 Q3 Q4 R P Y VX VY VZ WX WY WZ P0 P1 P2 P3 + obs_lower_bound = np.array([[-np.inf, -np.inf, 0., -1., -1., -1., -1., -np.pi, -np.pi, -np.pi, -np.inf, -np.inf, -np.inf, -np.inf, -np.inf, -np.inf, 0., 0., 0., 0.] for i in range(self.NUM_DRONES)]) + obs_upper_bound = np.array([[np.inf, np.inf, np.inf, 1., 1., 1., 1., np.pi, np.pi, np.pi, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, self.MAX_RPM, self.MAX_RPM, self.MAX_RPM, self.MAX_RPM] for i in range(self.NUM_DRONES)]) + return spaces.Box(low=obs_lower_bound, high=obs_upper_bound, dtype=np.float32) + + ################################################################################ + + def _computeObs(self): + """Returns the current observation of the environment. + + For the value of the state, see the implementation of `_getDroneStateVector()`. + + Returns + ------- + ndarray + An ndarray of shape (NUM_DRONES, 20) with the state of each drone. + + """ + return np.array([self._getDroneStateVector(i) for i in range(self.NUM_DRONES)]) + + ################################################################################ + + def _preprocessAction(self, + action + ): + """Pre-processes the action passed to `.step()` into motors' RPMs. + + Clips and converts a dictionary into a 2D array. + + Parameters + ---------- + action : ndarray + The (unbounded) input action for each drone, to be translated into feasible RPMs. + + Returns + ------- + ndarray + (NUM_DRONES, 4)-shaped array of ints containing to clipped RPMs + commanded to the 4 motors of each drone. + + """ + remapped_input = np.array([[ + action[i][2], # 0 + action[i][1], # 1 + action[i][3], # 2 + action[i][0] # 3 + ] for i in range(self.NUM_DRONES)]) # Betaflight SITL motor mapping + + ret = np.array(np.sqrt(self.MAX_THRUST / 4 / self.KF * remapped_input)) + assert(ret.shape == (self.NUM_DRONES, 4)), "Error in preprocess action" + return ret + + ################################################################################ + + def _computeReward(self): + """Computes the current reward value(s). + + Unused as this subclass is not meant for reinforcement learning. + + Returns + ------- + int + Dummy value. + + """ + return -1 + + ################################################################################ + + def _computeTerminated(self): + """Computes the current terminated value(s). + + Unused as this subclass is not meant for reinforcement learning. + + Returns + ------- + bool + Dummy value. + + """ + return False + + ################################################################################ + + def _computeTruncated(self): + """Computes the current truncated value(s). + + Unused as this subclass is not meant for reinforcement learning. + + Returns + ------- + bool + Dummy value. + + """ + return False + + ################################################################################ + + def _computeInfo(self): + """Computes the current info dict(s). + + Unused as this subclass is not meant for reinforcement learning. + + Returns + ------- + dict[str, int] + Dummy value. + + """ + return {"answer": 42} #### Calculated by the Deep Thought supercomputer in 7.5M years diff --git a/gym_pybullet_drones/examples/beta2.py b/gym_pybullet_drones/examples/beta2.py new file mode 100644 index 000000000..07384ba02 --- /dev/null +++ b/gym_pybullet_drones/examples/beta2.py @@ -0,0 +1,176 @@ +"""Control + Betaflight. + +Setup +----- +Step 1: Clone betaflight: + $ git clone https://github.com/betaflight/betaflight + $ cd betaflight/ + +Step 2: Comment out line `delayMicroseconds_real(50); // max rate 20kHz` + (https://github.com/betaflight/betaflight/blob/master/src/main/main.c#L52) + from Betaflight's `SIMULATOR_BUILD` and compile: + $ cd betaflight/ + $ make arm_sdk_install + $ make TARGET=SITL + +Step 3: Install betaflight configurator (https://github.com/betaflight/betaflight-configurator/releases): + $ wget https://github.com/betaflight/betaflight-configurator/releases/download/10.9.0/betaflight-configurator_10.9.0_amd64.deb + $ sudo dpkg -i betaflight-configurator_10.9.0_amd64.deb + If needed, also run: + $ apt install libgconf-2-4 + $ apt --fix-broken install + +Step 4: Load the configuration file onto the target using the BF configurator: + First, start the SITL controller: + $ ./obj/main/betaflight_SITL.elf + Type address `tcp://localhost:5761` (top right) and click `Connect` + Find the `Presets` tab (on the left) -> `Load backup` -> select file `../assets/beta.txt` + Restart `betaflight_SITL.elf` + +Example +------- +In one terminal run the SITL Betaflight: + + $ cd betaflight/ + $ ./obj/main/betaflight_SITL.elf + +In a separate terminal, run: + + $ cd gym-pybullet-drones/gym_pybullet_drones/examples/ + $ python beta.py + +""" +import time +import argparse +import numpy as np +import csv + +from transforms3d.quaternions import rotate_vector, qconjugate, mat2quat, qmult +from transforms3d.utils import normalized_vector + +from gym_pybullet_drones.utils.enums import DroneModel, Physics +from gym_pybullet_drones.envs.BetaAviary import BetaAviary +from gym_pybullet_drones.control.CTBRControl import CTBRControl +from gym_pybullet_drones.utils.Logger import Logger +from gym_pybullet_drones.utils.utils import sync, str2bool + +DEFAULT_DRONES = DroneModel("racer") +DEFAULT_PHYSICS = Physics("pyb") +DEFAULT_GUI = True +DEFAULT_PLOT = True +DEFAULT_USER_DEBUG_GUI = False +DEFAULT_SIMULATION_FREQ_HZ = 500 +DEFAULT_CONTROL_FREQ_HZ = 500 +DEFAULT_DURATION_SEC = 20 +DEFAULT_OUTPUT_FOLDER = 'results' + +def run( + drone=DEFAULT_DRONES, + physics=DEFAULT_PHYSICS, + gui=DEFAULT_GUI, + plot=DEFAULT_PLOT, + user_debug_gui=DEFAULT_USER_DEBUG_GUI, + simulation_freq_hz=DEFAULT_SIMULATION_FREQ_HZ, + control_freq_hz=DEFAULT_CONTROL_FREQ_HZ, + duration_sec=DEFAULT_DURATION_SEC, + output_folder=DEFAULT_OUTPUT_FOLDER, + ): + #### Create the environment with or without video capture ## + env = BetaAviary(drone_model=drone, + num_drones=1, + initial_xyzs=np.array([[.0, .0, .1]]), + initial_rpys=np.array([[.0, .0, .0]]), + physics=physics, + pyb_freq=simulation_freq_hz, + ctrl_freq=control_freq_hz, + gui=gui, + user_debug_gui=user_debug_gui + ) + + ctrl = CTBRControl(drone_model=drone) + + #### Obtain the PyBullet Client ID from the environment #### + PYB_CLIENT = env.getPyBulletClient() + + #### Initialize the logger ################################# + logger = Logger(logging_freq_hz=control_freq_hz, + num_drones=1, + output_folder=output_folder, + ) + + #### Run the simulation #################################### + with open("../assets/beta.csv", mode='r') as csv_file: + csv_reader = csv.DictReader(csv_file) + trajectory = iter([{ + "pos": np.array([ + float(row["p_x"]), + float(row["p_y"]), + float(row["p_z"]), + ]), + "vel": np.array([ + float(row["v_x"]), + float(row["v_y"]), + float(row["v_z"]), + ]), + } for row in csv_reader]) + action = np.zeros((1,4)) + ARM_TIME = 1. + TRAJ_TIME = 1.5 + START = time.time() + for i in range(0, int(duration_sec*env.CTRL_FREQ)): + t = i/env.CTRL_FREQ + #### Step the simulation ################################### + obs, reward, terminated, truncated, info = env.step(action, i) + + if t > env.TRAJ_TIME: + try: + target = next(trajectory) + action[0,:] = ctrl.computeControlFromState(control_timestep=env.CTRL_TIMESTEP, + state=obs[0], + target_pos=target["pos"], + target_vel=target["vel"] + ) + except: + break + + + #### Log the simulation #################################### + logger.log(drone=0, + timestamp=i/env.CTRL_FREQ, + state=obs[0] + ) + + #### Printout ############################################## + env.render() + + #### Sync the simulation ################################### + if gui: + pass + sync(i, START, env.CTRL_TIMESTEP) + + #### Close the environment ################################# + env.close() + + #### Save the simulation results ########################### + logger.save() + logger.save_as_csv("beta") # Optional CSV save + + #### Plot the simulation results ########################### + if plot: + logger.plot() + +if __name__ == "__main__": + #### Define and parse (optional) arguments for the script ## + parser = argparse.ArgumentParser(description='Test flight script using SITL Betaflight') + parser.add_argument('--drone', default=DEFAULT_DRONES, type=DroneModel, help='Drone model (default: BETA)', metavar='', choices=DroneModel) + parser.add_argument('--physics', default=DEFAULT_PHYSICS, type=Physics, help='Physics updates (default: PYB)', metavar='', choices=Physics) + parser.add_argument('--gui', default=DEFAULT_GUI, type=str2bool, help='Whether to use PyBullet GUI (default: True)', metavar='') + parser.add_argument('--plot', default=DEFAULT_PLOT, type=str2bool, help='Whether to plot the simulation results (default: True)', metavar='') + parser.add_argument('--user_debug_gui', default=DEFAULT_USER_DEBUG_GUI, type=str2bool, help='Whether to add debug lines and parameters to the GUI (default: False)', metavar='') + parser.add_argument('--simulation_freq_hz', default=DEFAULT_SIMULATION_FREQ_HZ, type=int, help='Simulation frequency in Hz (default: 240)', metavar='') + parser.add_argument('--control_freq_hz', default=DEFAULT_CONTROL_FREQ_HZ, type=int, help='Control frequency in Hz (default: 48)', metavar='') + parser.add_argument('--duration_sec', default=DEFAULT_DURATION_SEC, type=int, help='Duration of the simulation in seconds (default: 5)', metavar='') + parser.add_argument('--output_folder', default=DEFAULT_OUTPUT_FOLDER, type=str, help='Folder where to save logs (default: "results")', metavar='') + ARGS = parser.parse_args() + + run(**vars(ARGS))