made notebook run_sim signature same as optimize.run_sim()

%% Cell type:markdown id:d001f534 tags:

#### Notebook setup

%% Cell type:code id:be76bd53 tags:

``` python

%pip install -e .

```

%% Cell type:code id:5a776232 tags:

``` python

%reload_ext autoreload

%autoreload 2

%matplotlib inline

import time

import warnings

import os, sys

from copy import deepcopy

from types import SimpleNamespace

from pprint import pprint as print

import matplotlib.pyplot as plt

import gym

import numpy as np

from tqdm.auto import tqdm, trange

from multirotor.helpers import control_allocation_matrix, DataLog

from multirotor.vehicle import MotorParams, VehicleParams, PropellerParams, SimulationParams, BatteryParams

from multirotor.controller import (

    PosController, VelController,

    AttController, RateController,

    AltController, AltRateController,

    Controller

)

from multirotor.simulation import Multirotor, Propeller, Motor, Battery

from multirotor.coords import body_to_inertial, inertial_to_body, direction_cosine_matrix, euler_to_angular_rate

from multirotor.env import SpeedsMultirotorEnv, DynamicsMultirotorEnv

from multirotor.trajectories import Trajectory

from multirotor.visualize import plot_datalog

```

%% Cell type:code id:421b284b tags:

``` python

# Plotting/display parameters

# https://stackoverflow.com/a/21009774/4591810

float_formatter = "{:.3f}".format

np.set_printoptions(formatter={'float_kind':float_formatter})

SMALL_SIZE = 16

MEDIUM_SIZE = 16

BIGGER_SIZE = 20

plt.rc('font', size=SMALL_SIZE)          # controls default text sizes

plt.rc('axes', titlesize=MEDIUM_SIZE)     # fontsize of the axes title

plt.rc('axes', labelsize=BIGGER_SIZE, titlesize=BIGGER_SIZE)    # fontsize of the x and y labels

plt.rc('xtick', labelsize=MEDIUM_SIZE)    # fontsize of the tick labels

plt.rc('ytick', labelsize=MEDIUM_SIZE)    # fontsize of the tick labels

plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize

plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

```

%% Cell type:markdown id:9a31149f tags:

### Parameters

%% Cell type:code id:1ba2a2fe tags:

``` python

# Tarot T18 params

bp = BatteryParams(max_voltage=22.2)

mp = MotorParams(

    moment_of_inertia=5e-5,

    # resistance=0.27,

    resistance=0.081,

    k_emf=0.0265,

    # k_motor=0.0932,

    speed_voltage_scaling= 0.0347,

    max_current=38.

)

pp = PropellerParams(

    moment_of_inertia=1.86e-6,

    use_thrust_constant=True,

    k_thrust=9.8419e-05, # 18-inch propeller

    # k_thrust=5.28847e-05, # 15 inch propeller

    k_drag=1.8503e-06, # 18-inch propeller

    # k_drag=1.34545e-06, # 15-inch propeller

    motor=mp

)

vp = VehicleParams(

    propellers=[pp] * 8,

    battery=bp,

    # angles in 45 deg increments, rotated to align with

    # model setup in gazebo sim (not part of this repo)

    angles=np.linspace(0, -2*np.pi, num=8, endpoint=False) + 0.375 * np.pi,

    distances=np.ones(8) * 0.635,

    clockwise=[-1,1,-1,1,-1,1,-1,1],

    mass=10.66,

    inertia_matrix=np.asarray([

        [0.2206, 0, 0],

        [0, 0.2206, 0.],

        [0, 0, 0.4238]

    ])

)

sp = SimulationParams(dt=0.01, g=9.81)

```

%% Cell type:markdown id:48866609 tags:

### Multirotor

%% Cell type:markdown id:17f3b34d tags:

Simulating individual components of the multirotor. These make up the final `Multirotor` object.

%% Cell type:markdown id:0505e4f6 tags:

#### Motor

%% Cell type:code id:b3128c10 tags:

``` python

%matplotlib inline

# Plot motor speeds as a function of time and input voltage signal

plt.figure(figsize=(8,8))

motor = Motor(mp, sp)

for vsignal in [2, 4, 6, 8, 10, 12, 14, 16, 18, 20]:

    speeds = []

    motor.reset()

    speed = vsignal / mp.speed_voltage_scaling

    for i in range(200):

        speeds.append(motor.step(speed))

    plt.plot(speeds, label='%d rad/s' % speed)

plt.legend(ncol=2)

plt.ylabel('Speed rad/s')

plt.xlabel('Time /ms')

```

%% Cell type:markdown id:f3ddb119 tags:

Learning a linear relationship for the equation $V = k_{scaling} * speed$ for motors. This is useful for `SpeedsMultirotorEnv` which takes speed signals as the input. This constant converts speeds to applied voltages. The default value in`MotorParams` is 1, meaning the actions are voltage signals.

%% Cell type:code id:42ebfb97 tags:

``` python

from multirotor.helpers import learn_speed_voltage_scaling

def make_motor_fn(params, sp):

    from copy import deepcopy

    params = deepcopy(params)

    params.speed_voltage_scaling = 1.

    def motor_step(signal):

        m = Motor(params, sp)

        for i in range(100):

            s = m.step(signal)

        return s

    return motor_step

print('Voltage = %.5f * speed' % (learn_speed_voltage_scaling(make_motor_fn(mp, sp))))

```

%% Cell type:markdown id:8d42fff7 tags:

#### Propeller

%% Cell type:markdown id:6fb2f5c9 tags:

The propeller can use a numerically solved thrust relationship, where thrust depends on airspeed. Or the easier option of using thrust coefficient is available.

%% Cell type:code id:1669687b tags:

``` python

%matplotlib inline

# Plot propeller speed by numerically solving the thrust equation,

# *if* accurate propeller measurements are given in params

pp_ = deepcopy(pp)

pp_.use_thrust_constant = False # Set to true to just use k_thrust

prop = Propeller(pp_, sp)

plt.figure(figsize=(8,8))

speeds = np.linspace(0, 600, num=100)

for a in np.linspace(0, 10, 10, endpoint=False):

    thrusts = []

    for s in speeds:

        thrusts.append(prop.thrust(s, np.asarray([0, 0, a])))

    plt.plot(speeds, thrusts, label='%.1f m/s' % a)

plt.xlabel('Speed rad/s')

plt.ylabel('Thrust /N')

plt.title('Thrust with airspeed')

plt.legend(ncol=2)

```

%% Cell type:markdown id:e7bb489d tags:

#### Vehicle

%% Cell type:markdown id:2936761d tags:

Create a `Multirotor` object, given `VehicleParams` and `SimulationParams`

%% Cell type:code id:ecc23e14 tags:

``` python

# Combine propeller/motor/vehicle to get vehicle.

# Take off simulation

m = Multirotor(vp, sp)

log = DataLog(vehicle=m) # convenient logging class

m.reset()

m.state *= 0 # set to zero, reset() sets random values

action = m.allocate_control( # In this case action is allocated speed signals

    thrust=m.weight * 1.1,

    torques=np.asarray([0, 0, 0])

)

for i in range(500):

    m.step_speeds(action)

    log.log()

log.done_logging()

plt.plot(log.z)

```

%% Cell type:markdown id:90475ed5 tags:

### Gym Environment

%% Cell type:code id:857201b9 tags:

``` python

# this env takes the vector of [force_x, force_y, force_z, torque_x, torque_y, torque_z] to move

# the multirotor

env = DynamicsMultirotorEnv(Multirotor(vp, sp), max_rads=600)

env.reset()

log = DataLog(vehicle=env.vehicle)

for _ in range(100):

    env.step(np.asarray([0,0,env.vehicle.weight * 1.2, 0,0,0]))

    log.log()

log.done_logging()

plt.plot(log.z)

```

%% Cell type:code id:014b7077 tags:

``` python

# this env takes the vector of speed signals to move

# the multirotor

env = SpeedsMultirotorEnv(Multirotor(vp, sp))

env.reset()

log = DataLog(vehicle=env.vehicle)

for _ in range(100):

    env.step(np.ones(8) * 400)

    log.log()

log.done_logging()

plt.plot(log.z)

```

%% Cell type:markdown id:2c6535f4 tags:

### PID Controller

%% Cell type:markdown id:37d265bb tags:

This section explains how a PID controller is constructed. This is a cascaded PID architecture. See `Controller` docs

for more details.

%% Cell type:code id:33984887 tags:

``` python

# From PID parameters file

def get_controller(m: Multirotor, max_velocity=5., max_acceleration=3.):

    assert m.simulation.dt <= 0.1, 'Simulation time step too large.'

    pos = PosController(

        1.0, 0., 0., 1., vehicle=m,

        max_velocity=max_velocity, max_acceleration=max_acceleration,

        square_root_scaling=False, leashing=False

    )

    vel = VelController(

        2.0, 1.0, 0.5,

        max_err_i=max_acceleration,

        max_tilt=np.pi/12,

        vehicle=m)

    att = AttController(

        [2.6875, 4.5, 4.5],

        0, 0.,

        max_err_i=1.,

        vehicle=m)

    rat = RateController(

        [4., 4., 4.],

        0, 0,

        max_err_i=0.5,

        max_acceleration=1.,

        vehicle=m)

    alt = AltController(

        1, 0, 0,

        max_err_i=1, vehicle=m,

        max_velocity=max_velocity)

    alt_rate = AltRateController(

        5, 0, 0,

        max_err_i=1, vehicle=m)

    ctrl = Controller(

        pos, vel, att, rat, alt, alt_rate,

        period_p=0.1, period_a=0.01, period_z=0.1

    )

    return ctrl

```

%% Cell type:code id:d21d73fe tags:

``` python

%matplotlib inline

m = Multirotor(vp, sp)

ctrl = get_controller(m)

log = DataLog(vehicle=m, controller=ctrl)

for i in range(500):

    action = ctrl.step((0.01,0.1,1,0))

    # no allocation or motor simulation, for which we first need to

    # m.step_speeds(m.allocate_control(action[0], action[3:])

    # Instead, requested dynamics are fulfilled:

    dynamics = np.zeros(6, m.dtype)

    dynamics[2] = action[0]

    dynamics[3:] = action[1:]

    m.step_dynamics(dynamics)

    log.log()

log.done_logging()

plt.plot(log.actions[:,0], ls=':', label='thrust')

lines = plt.gca().lines

plt.twinx()

for s, axis in zip(log.actions.T[1:], ('x','y','z')):

    plt.plot(s, label=axis + '-torque')

plt.legend(handles=plt.gca().lines + lines)

```

%% Cell type:markdown id:cfbc9c25 tags:

#### Attitude Angle Controller

%% Cell type:code id:f012e1f8 tags:

``` python

m = Multirotor(vp, sp)

fz = m.weight

ctrl = get_controller(m)

att =  ctrl.ctrl_a

log = DataLog(vehicle=m, controller=ctrl, other_vars=('err',))

for i in range(5000):

    ref = np.asarray([np.pi/18, 0, 0])

    # action is prescribed euler rate

    action = att.step(ref, m.orientation, dt=sp.dt)

    # action = np.clip(action, a_min=-0.1, a_max=0.1)

    m.step_dynamics(np.asarray([0, 0, 0, *action]))

    log.log(err=att.err_p[0])

    log._actions[-1] = action

log.done_logging()

plt.plot(log.roll * 180 / np.pi)

plt.twinx()

plt.plot(log.actions[:,0], ls=':', label='Rate rad/s')

```

%% Cell type:markdown id:fdbd88c1 tags:

#### Attitude Rate Controller

%% Cell type:code id:9cc3b317 tags:

``` python

m = Multirotor(vp, sp)

fz = m.weight

ctrl = get_controller(m)

rat = ctrl.ctrl_r

att = ctrl.ctrl_a

log = DataLog(vehicle=m, controller=ctrl, other_vars=('err',))

for i in range(200):

    ref = np.asarray([np.pi/18, np.pi/12, 0])

    rate = att.step(ref, m.orientation, m.simulation.dt)

    torque = rat.step(rate, m.euler_rate, m.simulation.dt)

    action = np.clip(torque, a_min=-0.1, a_max=0.1)

    m.step_dynamics(np.asarray([0, 0, 0, *action]))

    log.log(err=rat.err_p[0])

    log._actions[-1] = action

log.done_logging()

plt.plot(log.roll * 180 / np.pi, c='r', label='roll')

plt.plot(log.pitch * 180 / np.pi, c='g', label='pitch')

plt.plot(log.yaw * 180 / np.pi, c='b', label='yaw')

plt.ylabel('Orientation /deg')

plt.legend()

plt.twinx()

plt.plot(log.actions[:,0], ls=':', c='r')

plt.plot(log.actions[:,1], ls=':', c='g')

plt.plot(log.actions[:,2], ls=':', c='b')

plt.ylabel('Torque / Nm')

plt.title('Ref orientation' + str(ref))

```

%% Cell type:markdown id:6370aa80 tags:

#### Altitude Controller

%% Cell type:code id:5fcd0d24 tags:

``` python

m = Multirotor(vp, sp)

ctrl = get_controller(m)

alt = ctrl.ctrl_z

alt_rate = ctrl.ctrl_vz

log = DataLog(vehicle=m, other_vars=('thrust',))

for i in range(5000):

    ref = np.asarray([1.])

    rate = alt.step(ref, m.position[2:], dt=0.1)

    action = alt_rate.step(rate, m.inertial_velocity[2:], dt=0.1)

    action = np.clip(action, a_min=-2*m.weight, a_max=2*m.weight)

    m.step_dynamics(np.asarray([0, 0, action[0], 0,0,0]))

    log.log(thrust=action)

    #log._actions[-1] = action

log.done_logging()

l = plt.plot(log.thrust.squeeze(), label='Thrust')

plt.twinx()

plt.plot(log.z, ls=':', label='Altitude /m')

plt.legend(handles=l+plt.gca().lines)

```

%% Cell type:markdown id:f4278c17 tags:

#### Position Controller

%% Cell type:code id:91ae8919 tags:

``` python

m = Multirotor(vp, sp)

ctrl = get_controller(m)

pos = ctrl.ctrl_p

vel = ctrl.ctrl_v

rat = ctrl.ctrl_r

att = ctrl.ctrl_a

log = DataLog(vehicle=m, other_vars=('err', 'torques'))

for i in range(100):

    ref = np.asarray([1.,0.])

    # converting position -> velocity -> angles

    velocity = pos.step(ref, m.position[:2], dt=0.1)

    angles = vel.step(velocity, m.velocity[:2], dt=0.1)[::-1]

    # attitude controller operates at higher frequency

    rate = att.step(np.asarray([*angles, 0]), m.orientation, dt=0.01)

    action = rat.step(rate, m.euler_rate, dt=0.01)

    # clipping torques to prevent over-reactions

    action = np.clip(action, a_min=-0.1, a_max=0.1)

    m.step_dynamics(np.asarray([0, 0, m.weight, *action]))

    log.log(err=pos.err[0], torques=action)

log.done_logging()

plt.plot(log.x, label='x')

plt.plot(log.err, label='x-err')

plt.ylabel('x /m')

l = plt.gca().lines

plt.twinx()

plt.plot(log.torques[:,1], ls=':', label='y-torque', c='c')

plt.plot(log.pitch * 180 / np.pi, ls='-.', label='Pitch', c='m')

plt.legend(handles=plt.gca().lines+l)

# plt.plot(log.actions[:,0] * 180 / np.pi, ls=':')

```

%% Cell type:markdown id:565cf2c9 tags:

### Parameter search

%% Cell type:markdown id:dc8f3d0d tags:

Using `optuna` to search over the space of PID controller parameters.

%% Cell type:code id:607c8b66 tags:

``` python

from multirotor.optimize import optimize, DEFAULTS

print(DEFAULTS)

```

%% Cell type:code id:f3b82f8d tags:

``` python

# search over parameter space

study = optimize(vp, sp, ntrials=100)

```

%% Cell type:code id:c1a374a6 tags:

``` python

# apply best parameters from study to controller, and run a simulation

from multirotor.optimize import run_sim, apply_params

env = DynamicsMultirotorEnv(Multirotor(vp, sp))

traj = Trajectory(env.vehicle, [[0,0,0]], proximity=1)

ctrl = get_controller(env.vehicle)

ctrl.set_params(**apply_params(None, params=study.best_params))

env.reset()

ctrl.reset()

log = run_sim(env, traj, ctrl)

plot_datalog(log)

```

%% Cell type:markdown id:9eb34672 tags:

### Simulation

%% Cell type:markdown id:e17df085 tags:

Combining `Multiotor` and `Controller` to run a simulation. First, defining waypoints:

%% Cell type:code id:b60ef8d3 tags:

``` python

# NASA flight test

# wp = np.asarray([

#     [0.0, 0.0, 30.0],

#     [164.0146725649829, -0.019177722744643688, 30.0],

#     [165.6418055187678, 111.5351051245816, 30.0],

#     [127.3337449710234, 165.73576059611514, 30.0],

#     [-187.28170707810204, 170.33217775914818, 45.0],

#     [-192.03130502498243, 106.30660058604553, 45.0],

#     [115.89920266153058, 100.8644210617058, 30.0],

#     [114.81859536317643, 26.80923518165946, 30.0],

#     [-21.459931490011513, 32.60508110653609, 30.0]

# ])

wp = np.asarray([

    [0,10,0],

    [10,10,0],

    [10,0,0],

    [0,0,0]

])

```

%% Cell type:markdown id:b43dae3d tags:

Then, defining a disturbance (for example, wind). The disturabance function takes time, `Multirotor`, and returns the forces in the *body frame* of the vehicle.

%% Cell type:code id:f5f52df4 tags:

``` python

def wind(t, m):