Commit 9a330022 authored by Ibrahim's avatar Ibrahim
Browse files

Added option for disturbances in gym env; 

Updated env step() method to return tuple as per updated Gym spec (5 elements instead of 4, new one being a 'truncate' variable.)
parent 237831b7

Demo.ipynb

+4 −2
Original line number Diff line number Diff line 
%% 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



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, ControllerParams

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, angular_to_euler_rate

from multirotor.env import SpeedsMultirotorEnv as LocalOctorotor

from multirotor.trajectories import Trajectory, GuidedTrajectory

```



%% 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

mp = MotorParams(

    moment_of_inertia=5e-5,

    resistance=0.27,

    k_emf=0.0265,

    speed_voltage_scaling=0.0347

)

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,

    # 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.001, g=9.81)

```



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



### Multirotor



%% 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 signal in [2, 4, 6, 8, 10, 12, 14, 16, 18, 20]:

    speeds = []

    motor.reset()

    for i in range(200):

        speeds.append(motor.step(signal))

    plt.plot(speeds, label='%dV' % signal)

plt.legend(ncol=2)

plt.ylabel('Speed rad/s')

plt.xlabel('Time /ms')

```



%% 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.

    m = Motor(params, sp)

    def motor_step(signal):

        for i in range(100):

            s = m.step(signal)

        return s

    return motor_step



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

```



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



#### Propeller



%% 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: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(

    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:2c6535f4 tags:



### PID Controller



%% Cell type:code id:33984887 tags:



``` python

# From PID parameters file

def get_controller(m: Multirotor):

    pos = PosController(

        1.0, 0., 0., 1., dt=1e-3, vehicle=m)

    vel = VelController(

        2.0, 1.0, 0.5, 1000., dt=1e-3, vehicle=m)

    att = AttController(

        [2.6875, 4.5, 4.5],

        0, 0.,

        1., dt=1e-3, vehicle=m)

    rat = RateController(

        [0.1655, 0.1655, 0.5],

        [0.135, 0.135, 0.018],

        [0.01234, 0.01234, 0.],

        [0.5,0.5,0.5], dt=1e-3, vehicle=m)

    alt = AltController(

        1, 0, 0,

        1, dt=1e-3, vehicle=m)

    alt_rate = AltRateController(

        5, 0, 0,

        1, dt=1e-3, vehicle=m)

    ctrl = Controller(

        pos, vel, att, rat, alt, alt_rate

    )

    return ctrl

```



%% Cell type:code id:d21d73fe tags:



``` python

%matplotlib inline

m = Multirotor(vp, sp)

ctrl = get_controller(m)

log = DataLog(controller=ctrl)

for i in range(100):

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

    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

att =  get_controller(m).ctrl_a

log = DataLog(vehicle=m, controller=att, 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)

    # 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=rat, other_vars=('err',))

for i in range(5000):

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

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

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

    action = np.clip(action, 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)

plt.twinx()

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

```



%% 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, controller=alt, other_vars=('thrust',))

for i in range(5000):

    ref = np.asarray([1.])

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

    action = alt_rate.step(rate, m.world_velocity[2:])

    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()



plt.plot(log.actions.squeeze())

plt.twinx()

plt.plot(log.z, ls=':')

```



%% 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, controller=pos, other_vars=('err', 'att_actions'))

for i in range(5000):

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

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

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

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

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

    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_p[0], att_actions=action)

log.done_logging()



plt.plot(log.position[:,0])

plt.plot(log.err)

# plt.plot(log.position[:,1])

plt.twinx()

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

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

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

```



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



### Simulation



%% Cell type:code id:f5f52df4 tags:



``` python

# def wind(t, m):

#     w_inertial = np.asarray([5 * np.sin(t * 2 * np.pi / 4000), 0, 0])

#     dcm = direction_cosine_matrix(*m.orientation)

#     return inertial_to_body(w_inertial, dcm)[:2]

wind = lambda t, m: [0, 0]

```



%% Cell type:code id:d60f1ca5 tags:



``` python

from pyscurve import plot_trajectory

```



%% Cell type:code id:a98de44f tags:



``` python

m = Multirotor(vp, sp)

env = LocalOctorotor(vehicle=m)



# for _ in range(5000):

#     state, *_ = env.step(np.asarray([0., 0., m.weight * 1.01, 0., 0., 0.]))



waypoints = [[0,50,2], [50,50,2], [50,0,2], [25,-25,2]]

traj = GuidedTrajectory(env.vehicle, waypoints, proximity=2)

# traj = Trajectory(env.vehicle, waypoints, proximity=2, resolution=None)



ctrl = get_controller(m)



errs = SimpleNamespace()

errs.pos = SimpleNamespace()

errs.pos.p, errs.pos.i, errs.pos.d = [], [], []

errs.att = SimpleNamespace()

errs.att.p, errs.att.i, errs.att.d = [], [], []

currents = []



log = DataLog(env.vehicle, ctrl,

              other_vars=('speeds','target', 'alloc_errs',

                          'ctrl_p', 'leash', 'err_p', 'ff_vel'))



for i, (pos, feed_forward_vel) in tqdm(enumerate(traj), leave=False, total=60000):

    if i==60000: break

    if i==1000: break

    # Generate reference for controller

    ref = np.asarray([*pos, 0.])

    # Get prescribed dynamics for system as thrust and torques

    dynamics = ctrl.step(ref, feed_forward_velocity=feed_forward_vel)

    thrust, torques = dynamics[0], dynamics[1:]

    # Allocate control: Convert dynamics into motor rad/s

    action = m.allocate_control(thrust, torques)

    # Send speeds to environment

    state, *_ = env.step(action)

    state, *_ = env.step(

        action, disturb_forces=0, disturb_torques=0

    )

    currents.append([p.motor.current for p in m.propellers])



    for (c, e) in zip((ctrl.ctrl_p, ctrl.ctrl_a), (errs.pos, errs.att)):

        e.p.append(c.err_p)

        e.i.append(c.err_i)

        e.d.append(c.err_d)

    alloc_errs = np.asarray([thrust, *torques]) - m.alloc @ action**2



    log.log(speeds=action, target=pos, alloc_errs=alloc_errs,

            ctrl_p=ctrl.ctrl_p.action, leash=ctrl.ctrl_p.leash,

            err_p=ctrl.ctrl_p.err, ff_vel=feed_forward_velocity)



    if np.any(np.abs(m.orientation[:2]) > np.pi/6): break



log.done_logging()

currents = np.asarray(currents)

```



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



``` python

plot_trajectory(traj.trajs[-1], dt=0.1)

```



%% Cell type:code id:81a8043e tags:



``` python

%matplotlib inline

plt.figure(figsize=(21,10.5))

plot_grid = (3,3)

plt.subplot(*plot_grid,1)



n = len(log)



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

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

plt.plot(log.y, label='y', c='g')

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

plt.plot(log.z, label='z', c='b')

lines = plt.gca().lines[::2]

plt.ylabel('Position /m')

plt.twinx()

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

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

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

plt.ylabel('Orientation /deg')

plt.legend(handles=plt.gca().lines + lines, ncol=2)

plt.title('Position and Orientation')



plt.subplot(*plot_grid,2)

for i in range(log.speeds.shape[1]):

    l, = plt.plot(log.speeds[:,i], label='prop %d' % i)

#     plt.plot(speeds[:,i], c=l.get_c())

lines = plt.gca().lines

plt.legend(handles=lines, ncol=2)

plt.title('Motor speeds /RPM')





plt.subplot(*plot_grid,3)

v_world = np.zeros_like(log.velocity)

for i, (v, o) in enumerate(zip(log.velocity, log.orientation)):

    dcm = direction_cosine_matrix(*o)

    v_world[i] = body_to_inertial(v, dcm)

for i, c, a in zip(range(3), 'rgb', 'xyz'):

    plt.plot(v_world[:,i], label='Velocity %s' % a, c=c)

#     plt.plot(velocities[:,i], label='Velocity %s' % a, c=c)

plt.legend()

plt.title('Velocities')



plt.subplot(*plot_grid,4)

plt.title('Controller allocated dynamics')

l = plt.plot(log.actions[:,0], label='Ctrl Thrust')

plt.ylabel('Force /N')

plt.twinx()

for i, c, a in zip(range(3), 'rgb', 'xyz'):

    plt.plot(log.actions[:,1+i], label='Ctrl Torque %s' % a, c=c)

plt.ylabel('Torque /Nm')

plt.legend(handles=plt.gca().lines + l, ncol=2)



plt.subplot(*plot_grid,5)

lines = plt.plot(log.alloc_errs[:, 0], label='Thrust err', c='b')

plt.ylabel('Thrust /N')

plt.twinx()

plt.plot(log.alloc_errs[:, 1], label='Torque x err', ls=':')

plt.plot(log.alloc_errs[:, 2], label='Torque y err', ls=':')

plt.plot(log.alloc_errs[:, 3], label='Torque z err', ls=':')

plt.legend(handles = plt.gca().lines + lines, ncol=2)

plt.ylabel('Torque /Nm')

plt.title('Allocation Errors')



plt.subplot(*plot_grid,6)

plt.plot(log.target[:,0], log.target[:,1], label='Prescribed traj')

plt.plot(log.x, log.y, label='Actual traj', ls=':')

plt.gca().set_aspect('equal', 'box')

plt.title('XY positions /m')

plt.legend()



plt.tight_layout()

```



%% Cell type:code id:c17d8b15 tags:



``` python

%matplotlib inline

for e in (errs.pos, errs.att):

    e.p = np.asarray(e.p)

    e.i = np.asarray(e.i)

    e.d = np.asarray(e.d)

lines = []

plt.figure(figsize=(21,24))

plt.subplot(6,1,1)

plt.plot(errs.pos.p[:,0], label='Pos-x P', c='r', ls='-')

plt.plot(errs.pos.p[:,1], label='Pos-y P', c='g', ls='-')

lines += plt.gca().lines

plt.title('Position P errors')

plt.subplot(6,1, 2)

plt.plot(errs.pos.i[:, 0], label='Pos-x I', c='r', ls=':')

plt.plot(errs.pos.i[:, 1], label='Pos-y I', c='g', ls=':')

lines += plt.gca().lines

plt.title('Position I errors')

plt.subplot(6,1, 3)

plt.plot(errs.pos.d[:, 0], label='Pos-x D', c='r', ls='-.')

plt.plot(errs.pos.d[:, 1], label='Pos-y D', c='g', ls='-.')

lines += plt.gca().lines

plt.legend(handles=lines)

plt.title('Position D errors')

lines = []

plt.subplot(6,1,4)

plt.plot(errs.att.p[:,0], label='Att-x P', c='r', ls='-')

plt.plot(errs.att.p[:,1], label='Att-y P', c='g', ls='-')

lines += plt.gca().lines

plt.title('Attitude P errors')

plt.subplot(6,1, 5)

plt.plot(errs.att.i[:, 0], label='Att-x I', c='r', ls=':')

plt.plot(errs.att.i[:, 1], label='Att-y I', c='g', ls=':')

lines += plt.gca().lines

plt.title('Attitude I errors')

plt.subplot(6,1, 6)

plt.plot(errs.att.d[:, 0], label='Att-x D', c='r', ls='-.')

plt.plot(errs.att.d[:, 1], label='Att-y D', c='g', ls='-.')

lines += plt.gca().lines

plt.legend(handles=lines)

plt.title('Attitude D errors')

```



%% Cell type:code id:d9d99572 tags:



``` python

%matplotlib notebook

fig = plt.figure()

xlim = ylim = zlim = (np.min(log.position), np.max(log.position))

ax = fig.add_subplot(projection='3d', xlim=xlim, ylim=ylim, zlim=zlim)

ax.set_xlabel('x')

ax.set_ylabel('y')

ax.set_zlabel('z')

ax.plot(log.x, log.y, log.z)

```



%% Cell type:markdown id:e87063b0 tags:



### Variations



%% Cell type:code id:7c0024dc tags:



``` python

from multirotor.helpers import vehicle_params_factory

```



%% Cell type:code id:4a591d5b tags:



``` python

vp4 = vehicle_params_factory(4, m_prop=0.125, d_prop=0.25, params=pp, m_body=4)

vp6 = vehicle_params_factory(6, m_prop=0.125, d_prop=0.5, params=pp, m_body=6)

vp8 = vehicle_params_factory(8, m_prop=0.125, d_prop=0.75, params=pp, m_body=8)

wind = lambda t, m: [2, 0]

logs = []

for vpx in tqdm([vp4, vp6, vp8], leave=False):

    m = Multirotor(vpx, sp)

    env = LocalOctorotor(vehicle=m, allocate=True, max_rads=400)

    env.reset()

    traj = Trajectory([[2.5,5,0, 2000], [5,0,0,2000], [0,0,0,2000]],

                      env.vehicle, proximity=0.2, resolution=None)



    ctrl = Controller(

        # (Outer) Position controller has low sensitivity

        PosController(0.2, 0.02, 0., 1., dt=sp.dt, vehicle=m, max_tilt=np.pi/18),

        # (Inner) Attitude controller is more responsive (except for yaw, which we are not controlling)

        AttController(np.asarray([50., 50., 0.]),

                      np.asarray([50., 50., 0.]),

                      np.asarray([1., 1., 0.]), 1., dt=sp.dt, vehicle=env.vehicle),

        # Altitude controller is more responsive as well

        AltController(50, 2, 0, 1, dt=1e-3, vehicle=m)

    )



    log = DataLog(env.vehicle, ctrl, 'targets')



    for i, pos in tqdm(enumerate(traj), leave=False, total=60000):

        if i==60000: break

        # Get prescribed dynamics for system

        reference = np.asarray([*pos, 0.])

        dynamics = ctrl.step(reference)

        thrust, torques = dynamics[0], dynamics[1:]

        forces = np.asarray([0, 0, thrust])

        # Add disturbances

        f[:2] += wind(i, env.vehicle)

        # apply to simulation

        state, *_ = env.step(np.asarray([*forces, *torques]))



        log.log(targets=pos)



    log.done_logging()

    logs.append(log)

```



%% Cell type:code id:0b8e3a15 tags:



``` python

%matplotlib inline

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

for log, label, ls, lw in zip(

    logs,

    ('quad-rotor', 'hexa-rotor', 'octo-rotor'),

    ('-','--',':'),

    (2,3,4)

):

    plt.plot(log.x, log.y, label=label, ls=ls, lw=lw)

plt.gca().set_aspect('equal', 'box')

plt.title('XY positions /m')

plt.grid(which='both')

# plt.xlim(0, 5)

# plt.ylim(0, 5)

plt.legend()

```

multirotor/env.py

+53 −8
Original line number Diff line number Diff line
@@ -2,6 +2,7 @@ 
This module defines OpenAI Gym compatible classes based on the Multirotor class.

"""



from typing import Tuple

import numpy as np

import gym
@@ -59,21 +60,43 @@ class DynamicsMultirotorEnv(BaseMultirotorEnv): 
        self.max_rads = max_rads





    def step(self, action: np.ndarray):

    def step(

        self, action: np.ndarray, disturb_forces: np.ndarray=0.,

        disturb_torques: np.ndarray=0.

    ) -> Tuple[np.ndarray, float, bool, bool, dict]:

        """

        Step environment by providing dynamics acting in local frame.



        Parameters

        ----------

        action : np.ndarray

            An array of x,y,z forces and x,y,z torques in local frame.

        disturb_forces : np.ndarray, optional

            Disturbinng x,y,z forces in the vehicle's local frame, by default 0.

        disturb_torques : np.ndarray, optional

            Disturbing x,y,z torques in the vehicle's local frame, by default 0.



        Returns

        -------

        Tuple[np.ndarray, None, None, None]

            The state and other environment variables.

        """

        if self.allocate:

            speeds = self.vehicle.allocate_control(action[2], action[3:6])

            speeds = np.clip(speeds, a_min=0, a_max=self.max_rads)

            forces, torques = self.vehicle.get_forces_torques(speeds, self.vehicle.state)

            action = np.concatenate((forces, torques))

        action[:3] += disturb_forces

        action[3:] += disturb_torques

        self.vehicle.step_dynamics(u=action)

        return self.state, None, None, None

        return self.state, 0., False, False, {}







class SpeedsMultirotorEnv(BaseMultirotorEnv):





    def __init__(self, vehicle: Multirotor=None, max_rads=None) -> None:

    def __init__(self, vehicle: Multirotor=None) -> None:

        super().__init__(vehicle=vehicle)

        

        self.action_space = gym.spaces.Box(
@@ -82,10 +105,32 @@ class SpeedsMultirotorEnv(BaseMultirotorEnv): 
            dtype=np.float32,

            shape=(len(vehicle.propellers),)  # action for each propeller

        )

        self.max_rads = max_rads





    def step(self, action: np.ndarray):

        action = np.clip(action, a_min=0, a_max=self.max_rads)

        self.vehicle.step_speeds(u=action)

        return self.state, None, None, None
 
 No newline at end of file
 
    def step(

        self, action: np.ndarray, disturb_forces: np.ndarray=0.,

        disturb_torques: np.ndarray=0.

    ) -> Tuple[np.ndarray, float, bool, bool, dict]:

        """

        Step environment by providing speed signal.



        Parameters

        ----------

        action : np.ndarray

            An array of speed signals.

        disturb_forces : np.ndarray, optional

            Disturbinng x,y,z forces in the velicle's local frame, by default 0.

        disturb_torques : np.ndarray, optional

            Disturbing x,y,z torques in the vehicle's local frame, by default 0.



        Returns

        -------

        Tuple[np.ndarray, None, None, None]

            The state and other environment variables.

        """

        self.vehicle.step_speeds(

            u=action,

            disturb_forces=disturb_forces,

            disturb_torques=disturb_torques

        )

        return self.state, 0., False, False, {}
 
 No newline at end of file

multirotor/simulation.py

+19 −4
Original line number Diff line number Diff line
@@ -412,7 +412,10 @@ class Multirotor: 
        return np.around(xdot, self.dxdt_decimals)





    def dxdt_speeds(self, t: float, x: np.ndarray, u: np.ndarray, params=None):

    def dxdt_speeds(

        self, t: float, x: np.ndarray, u: np.ndarray,

        disturb_forces: np.ndarray=0., disturb_torques: np.ndarray=0., params=None

    ):

        """

        Calculate the rate of change of state given the propeller speeds on the

        system (rad/s).
@@ -425,6 +428,10 @@ class Multirotor: 
            State of the vehicle.

        u : np.ndarray

            A p-vector of propeller speeds (rad/s), where p=number of propellers.

        disturb_forces : np.ndarray, optional

            Disturbinng x,y,z forces in the vehicle's local frame, by default 0.

        disturb_torques : np.ndarray, optional

            Disturbing x,y,z torques in the vehicle's local frame, by default 0.



        Returns

        -------
@@ -437,7 +444,7 @@ class Multirotor: 
        forces, torques = self.get_forces_torques(

            u, x)

        xdot = apply_forces_torques(

            forces, torques, x, self.simulation.g,

            forces+disturb_forces, torques+disturb_torques, x, self.simulation.g,

            self.params.mass, self.params.inertia_matrix, self.params.inertia_matrix_inverse)

        return np.around(xdot, self.dxdt_decimals)
@@ -470,7 +477,10 @@ class Multirotor: 
        return self.state





    def step_speeds(self, u: np.ndarray) -> np.ndarray:

    def step_speeds(

        self, u: np.ndarray, disturb_forces: np.ndarray=0.,

        disturb_torques: np.ndarray=0.

    ) -> np.ndarray:

        """

        Given the n-vector of propeller speed signals, calculate

        the next state of the vehicle. Where n is number of propellers.
@@ -481,6 +491,10 @@ class Multirotor: 
            The speed signals to be sent to each propeller's step() method. Can

            be the actual speed (rad/s) or the voltage signal (V) if a motor

            is used and MotorParams.speed_voltage_scaling constant is set.

        disturb_forces : np.ndarray, optional

            Disturbinng x,y,z forces in the vehicle's local frame, by default 0.

        disturb_torques : np.ndarray, optional

            Disturbing x,y,z torques in the vehicle's local frame, by default 0.



        Returns

        -------
@@ -490,7 +504,8 @@ class Multirotor: 
        self.t += self.simulation.dt

        self._dxdt = self.dxdt_speeds(t=self.t, x=self.state, u=u)

        self.state = odeint(

            self.dxdt_speeds, self.state, (0, self.simulation.dt), args=(u,),

            self.dxdt_speeds, self.state, (0, self.simulation.dt),

            args=(u, disturb_forces, disturb_torques),

            rtol=1e-4, atol=1e-4, tfirst=True

        )[-1]

        self.state = np.around(self.state, 4)